def calc_tensor_sparta_yhat(config): # load all required inputs tmp_config = config['task'].copy() tmp_config.pop('dataset') tmp_config['dataset'] = tmp_config.pop('dataset_info') tmp_config['dataset']['input'] = ['T1', 'T2', 'T3', 'T4', 'inv1', 'inv2'] tmp_config['dataset']['output'] = 'bDelta' # tmp_config['dataset']['skip_wall'] = False RANSdata, _ = load_frozen_RANS_dataset(tmp_config['dataset']) T1 = RANSdata[:, 0] T2 = RANSdata[:, 1] T3 = RANSdata[:, 2] T4 = RANSdata[:, 3] inv1 = RANSdata[:, 4] inv2 = RANSdata[:, 5] model1 = False model2 = False model3 = True if model1: yhat_sparta = (24.94 * inv1**2 + 2.65 * inv2) * T1 + 2.96 * T2 + ( 2.49 * inv2 + 20.05) * T3 + (2.49 * inv1 + 14.93) * T4 if model2: yhat_sparta = T1*(0.46*inv1**2 + 11.68*inv2 - 0.30*(inv2**2) + 0.37) + T2*(-12.25*inv1 - 0.63*(inv2**2) + 8.23) + \ T3*(-1.36*inv2 - 2.44) + T4*(-1.36*inv1 + 0.41*inv2 - 6.52) if model3: yhat_sparta = T1*(0.11*inv1*inv2 + 0.27*inv1*(inv2**2) - 0.13*inv1*(inv2**3) + 0.07*inv1*(inv2**4) \ + 17.48*inv1 + 0.01*(inv1**2)*inv2 + 1.251*(inv1**2) + 3.67*inv2 + 7.52*(inv2**2) - 0.3) \ + T2*(0.17*inv1*(inv2**2) - 0.16*inv1*(inv2**2) - 36.25*inv1 - 2.39*(inv1**2) + 19.22*inv2 + 7.04) \ + T3*(-0.22*(inv1**2) + 1.8*inv2 + 0.07*(inv2**2)+2.65) + T4*(0.2*(inv1**2) - 5.23*inv2 - 2.93) return yhat_sparta
def contourplot_results_scalar(results, config): # re-read config file in output directory to find in and outputs logdir = config['training']['logdir'] with open(logdir + '/config.json', encoding='utf-8') as f: config = json.load(f) name = results['name'] seed = results['seed'] cases = ['PH10595', 'CD12600', 'CBFS13700'] for case in cases: tmp_config = config['task']['dataset'].copy() tmp_config['name'] = case X, y = load_frozen_RANS_dataset(tmp_config) # get yhat reduced field (used for training) yhat, _ = results['program'].execute(X) inv_nrmse = 1 / (1 + np.sqrt(np.mean((y - yhat)**2)) / np.std(y)) # reload data, now containing all points for plotting tmp_config['skip_wall'] = False X, y = load_frozen_RANS_dataset(tmp_config) # get yhat full field yhat, _ = results['program'].execute(X) # get mesh data frozen = pickle.load( open(f'turbulence/frozen_data/{case}_frozen_var.p', 'rb')) data_i = frozen['data_i'] mesh_x = data_i['meshRANS'][0, :, :] mesh_y = data_i['meshRANS'][1, :, :] filename = f'{logdir}/dsr_{name}_{seed}_contour_{case}' case_contourplots(mesh_x, mesh_y, y, yhat, filename, inv_nrmse, config['task']['dataset']['skip_wall'])
def retrospecitvely_plot_contours(logdir, with_sparta=True): with open(logdir + '/config.json', encoding='utf-8') as f: config = json.load(f) cases = ['PH10595', 'CD12600', 'CBFS13700'] for case in cases: config['task']['dataset']['name'] = case config['task']['dataset']['skip_wall'] = False X, y = load_frozen_RANS_dataset(config['task']['dataset']) ysparta = 2 * 1.4 * X[:, 0] * X[:, 1] frozen = pickle.load( open(f'turbulence/frozen_data/{case}_frozen_var.p', 'rb')) data_i = frozen['data_i'] mesh_x = data_i['meshRANS'][0, :, :] mesh_y = data_i['meshRANS'][1, :, :] # filename = f'{logdir}/dsr_{name}_{seed}_contour_{case}' # case_contourplots(mesh_x, mesh_y, y, yhat, filename) files = os.listdir(logdir) for filename in files: if filename[-7:] == 'hof.csv': df = pd.read_csv(f'{logdir}/{filename}') df_row = df.iloc[0] name = config['task']['name'] seed = [int(s) for s in filename.split('_') if s.isdigit()][0] filename = f'{logdir}/dsr_{name}_{seed}_contour_{case}' yhat = eval_expression(df_row['expression'], X) if with_sparta: case_contourplots_with_sparta(mesh_x, mesh_y, y, yhat, ysparta, filename) else: case_contourplots(mesh_x, mesh_y, y, yhat, filename) plt.close('all')
def bDelta_magnitudes(model_file): # # Possibly use this for all bDelta models! # df_models = pd.DataFrame() # for case in ['PH', 'CD', 'CBFS']: # file = f'../logs_completed/kDef_{case}/kDef_{case}_selected_models_CFD_results.csv' # # file = f'../logs_completed/bDel_{case}/bDel_{case}_selected_models_CFD_results_full_bDelta.csv' # df_case = pd.read_csv(file) # df_models = pd.concat([df_models, df_case], axis=0, ignore_index=True) # sort by training reward? df_models = pd.read_csv(model_file) df_models[['x1term','x2term','x3term','x4term']] = df_models.apply(lambda x: split_expression(x['batch_r_max_expression']), axis=1, result_type='expand') # Load the config file with open('/'.join(model_file.split('/')[:-1] + ['config.json']), encoding='utf-8') as f: config = json.load(f) # config['task']['dataset']['skip_wall'] = False X, y = load_frozen_RANS_dataset(config['task']['dataset']) vardict = {} for ii in range(X.shape[1]): vardict[f'x{ii+1}'] = X[:,ii] for ii in range(4): strings = df_models[f'x{ii+1}term'].values strings = [str.replace(f'x{ii+1}*', '') for str in strings] strings = [str.replace(f'*x{ii+1}', '') for str in strings] strings = [str.replace(f'x{ii+1}', '0') for str in strings] df_models[f'x{ii+1}term'] = strings df_models[[f'x{ii+1}corr', f'x{ii+1}absmean', f'x{ii+1}mean']] = df_models.apply(lambda x: tensor_expr_eval(x[f'x{ii+1}term'], vardict, y, f'x{ii+1}', x['batch_r_max_expression']), axis=1, result_type='expand') if len(df_models['training_case'].unique()) == 1: training_case = df_models['training_case'].unique()[0] if training_case == 'PH': sort_by_CFD = ['PH_nmse', 'CD_nmse', 'CBFS_nmse'] sort_by_r_max = ['r_max_PH', 'r_max_CD', 'r_max_CBFS'] rlabel = r'\; PH_{10595}$' # rlabel = r'$r_{max} \; PH_{10595}$' elif training_case == 'CD': sort_by_CFD = ['CD_nmse', 'PH_nmse', 'CBFS_nmse'] sort_by_r_max = ['r_max_CD', 'r_max_PH', 'r_max_CBFS'] rlabel = '\; CD_{12600}$' # rlabel = '$r_{max} \; CD_{12600}$' else: sort_by_CFD = ['CBFS_nmse', 'PH_nmse', 'CD_nmse'] sort_by_r_max = ['r_max_CBFS', 'r_max_PH', 'r_max_CD'] rlabel = '\; CBFS_{13700}$' df_sorted = df_models.sort_values(sort_by_r_max, ascending=[False, False, False], ignore_index=True) x = np.arange(df_sorted.shape[0]) + 1 markersize = 30 lw = 1.5 figsize = (24, 6) cm = 1 / 2.54 # centimeters in inches # plot r_max: plt.figure(figsize=tuple([val*cm for val in list(figsize)])) # plt.plot(x, df_sorted[f'r_max_{training_case}'] ) add_scatter(x, df_sorted, 'x1corr', 'C0', markersize, lw, r'$g^{(1)} T^{(1)} $', 'd') add_scatter(x, df_sorted, 'x2corr', 'C1', markersize, lw, r'$g^{(2)} T^{(2)} $', '^') add_scatter(x, df_sorted, 'x3corr', 'C2', markersize, lw, r'$g^{(3)} T^{(3)} $', 'v') add_scatter(x, df_sorted, 'x4corr', 'C3', markersize, lw, r'$g^{(4)} T^{(4)} $', 'o') plt.ylabel(r'$\rho_c$') plt.xlabel('Models') # plt.ylim([-18, 18]) plt.xticks(np.arange(0,100,10)) ax = plt.gca() ax.xaxis.grid(linestyle=':') ax.yaxis.grid(linestyle=':') plt.legend(prop={'size': 8}) ax2 = ax.twinx() ax2.set_ylabel(r'$r_{max} $') ax2.plot(x, df_sorted[f'r_max_{training_case}'] , label=r'$r_{max}' + rlabel, zorder=-1, c='black', linestyle=(0, (1, 1))) if training_case == 'CBFS': ax2.set_yticks([0.55, 0.60, 0.65]) lines_1, labels_1 = ax.get_legend_handles_labels() lines_2, labels_2 = ax2.get_legend_handles_labels() lines = lines_1 + lines_2 labels = labels_1 + labels_2 ax.legend(lines, labels, prop={'size': 10}, loc='upper center', bbox_to_anchor=(0.5, 1.23), ncol=5) plt.savefig(f'../logs_completed/aa_plots/bDelCorr_r_max_{training_case}.eps', format='eps', bbox_inches='tight') df_sorted = df_models.sort_values(sort_by_CFD, ascending=[True, True, True], ignore_index=True) # plot CFD errors: plt.figure(figsize=tuple([val*cm for val in list(figsize)])) # plt.plot(x, df_sorted[f'r_max_{training_case}'] ) add_scatter(x, df_sorted, 'x1corr', 'C0', markersize, lw, r'$g^{(1)} T^{(1)} $', 'd') add_scatter(x, df_sorted, 'x2corr', 'C1', markersize, lw, r'$g^{(2)} T^{(2)} $', '^') add_scatter(x, df_sorted, 'x3corr', 'C2', markersize, lw, r'$g^{(3)} T^{(3)} $', 'v') add_scatter(x, df_sorted, 'x4corr', 'C3', markersize, lw, r'$g^{(4)} T^{(4)} $', 'o') plt.ylabel(r'$\rho_c$') plt.xlabel('Models') # plt.ylim(ylim) plt.xticks(np.arange(0,100,10)) ax = plt.gca() ax.xaxis.grid(linestyle=':') ax.yaxis.grid(linestyle=':') plt.legend(prop={'size': 8}) ax2 = ax.twinx() ax2.set_ylabel(r'$\varepsilon (U) / \varepsilon(U_0)$') ax2.plot(x, df_sorted[f'{training_case}_nmse'] , label=r'$\varepsilon (U) / \varepsilon(U_0)' + rlabel, zorder=-1, c='black', linestyle=(0, (1, 1))) ax2.set_ylim([0, 4]) lines_1, labels_1 = ax.get_legend_handles_labels() lines_2, labels_2 = ax2.get_legend_handles_labels() lines = lines_1 + lines_2 labels = labels_1 + labels_2 ax.legend(lines, labels, prop={'size': 10}, loc='upper center', bbox_to_anchor=(0.5, 1.23), ncol=5) plt.savefig(f'../logs_completed/aa_plots/bDelCorr_CFD_err_{training_case}.eps', format='eps', bbox_inches='tight')
def make_regression_task(name, function_set, enforce_sum, dataset, dataset_info, metric="inv_nrmse", metric_params=(1.0,), extra_metric_test=None, extra_metric_test_params=(), reward_noise=0.0, reward_noise_type="r", threshold=1e-12, normalize_variance=False, protected=False): """ Factory function for regression rewards. This includes closures for a dataset and regression metric (e.g. inverse NRMSE). Also sets regression- specific metrics to be used by Programs. Parameters ---------- name : str or None Name of regression benchmark, if using benchmark dataset. function_set : list or None List of allowable functions. If None, uses function_set according to benchmark dataset. dataset : dict, str, or tuple If dict: .dataset.BenchmarkDataset kwargs. If str: filename of dataset. If tuple: (X, y) data metric : str Name of reward function metric to use. metric_params : list List of metric-specific parameters. extra_metric_test : str Name of extra function metric to use for testing. extra_metric_test_params : list List of metric-specific parameters for extra test metric. reward_noise : float Noise level to use when computing reward. reward_noise_type : "y_hat" or "r" "y_hat" : N(0, reward_noise * y_rms_train) is added to y_hat values. "r" : N(0, reward_noise) is added to r. normalize_variance : bool If True and reward_noise_type=="r", reward is multiplied by 1 / sqrt(1 + 12*reward_noise**2) (We assume r is U[0,1]). protected : bool Whether to use protected functions. threshold : float Threshold of NMSE on noiseless data used to determine success. Returns ------- task : Task Dynamically created Task object whose methods contains closures. """ X_test = y_test = y_test_noiseless = None # Benchmark dataset config if isinstance(dataset, dict): dataset["name"] = name benchmark = BenchmarkDataset(**dataset) X_train = benchmark.X_train y_train = benchmark.y_train X_test = benchmark.X_test y_test = benchmark.y_test y_test_noiseless = benchmark.y_test_noiseless # Unless specified, use the benchmark's default function_set if function_set is None: function_set = benchmark.function_set # Dataset filename elif isinstance(dataset, str): df = pd.read_csv(dataset, header=None) # Assuming data file does not have header rows X_train = df.values[:, :-1] y_train = df.values[:, -1] # sklearn-like (X, y) data elif isinstance(dataset, tuple): X_train = dataset[0] y_train = dataset[1] if X_test is None: X_test = X_train y_test = y_train y_test_noiseless = y_test X_train_full = X_train y_train_full = y_train if dataset_info['name'] in ['PH10595', 'CBFS13700', 'CD12600']: dummy_config = dataset_info.copy() dummy_config['name'] = 'PH10595' X_train_PH, y_train_PH = load_frozen_RANS_dataset(dummy_config) dummy_config = dataset_info.copy() dummy_config['name'] = 'CBFS13700' X_train_CBFS, y_train_CBFS = load_frozen_RANS_dataset(dummy_config) dummy_config = dataset_info.copy() dummy_config['name'] = 'CD12600' X_train_CD, y_train_CD = load_frozen_RANS_dataset(dummy_config) if function_set is None: print("WARNING: Function set not provided. Using default set.") function_set = ["add", "sub", "mul", "div", "sin", "cos", "exp", "log"] # Save time by only computing these once var_y_test = np.var(y_test) var_y_test_noiseless = np.var(y_test_noiseless) # Define closures for metric metric, invalid_reward, max_reward, ad_metric_traversal = make_regression_metric(metric, y_train, *metric_params) ad_metric_start_idx = len(ad_metric_traversal) - 1 if extra_metric_test is not None: print("Setting extra test metric to {}.".format(extra_metric_test)) metric_test, _, _ = make_regression_metric(extra_metric_test, y_test, *extra_metric_test_params) assert reward_noise >= 0.0, "Reward noise must be non-negative." if reward_noise: assert reward_noise_type in ["y_hat", "r"], "Reward noise type not recognized." rng = np.random.RandomState(0) y_rms_train = np.sqrt(np.mean(y_train ** 2)) if reward_noise_type == "y_hat": scale = reward_noise * y_rms_train elif reward_noise_type == "r": scale = reward_noise def data_shuffle(seed): nonlocal X_train_full, y_train_full idx = np.arange(y_train_full.shape[0]) np.random.seed(seed) idx = np.random.permutation(idx) X_train_full = X_train_full[idx, :] y_train_full = y_train_full[idx] def rotate_batch(batch_size, data_set=None, rotate=True): """ :param batch_size: int or None number of points in batch to be returned, if None full dataset is selected :param data_set: string or None if dataset in ['PH', 'CD', 'CBFS'] the X_train and y_train are changed to different dataset for evaluation :return: Nothing is returned, X_train and y_train are changed non locally """ nonlocal X_train, y_train, ad_metric_traversal, X_train_full, y_train_full if data_set == 'PH': # switch to PH nonlocal X_train_PH, y_train_PH X_train = X_train_PH y_train = y_train_PH elif data_set == 'CD': # switch to CD nonlocal X_train_CD, y_train_CD X_train = X_train_CD y_train = y_train_CD elif data_set == 'CBFS': # switch to CBFS nonlocal X_train_CBFS, y_train_CBFS X_train = X_train_CBFS y_train = y_train_CBFS else: # set up train batch if batch_size: if rotate: # rotate full dataset so next selection is of new points X_train_full = np.roll(X_train_full, -batch_size, axis=0) y_train_full = np.roll(y_train_full, -batch_size) X_train = X_train_full[:batch_size, :] y_train = y_train_full[:batch_size] ad_metric_traversal[-2] = AD_PlaceholderConstant(value=y_train, name='y') def reward(p): # Compute estimated values y_hat, invalid_indices = p.execute(X_train) p.invalid_tokens = invalid_indices # For invalid expressions, return invalid_reward # p.invalid_tokens = [token.invalid for token in p.traversal] if p.invalid: if invalid_indices is None: p.invalid = 1 else: p.invalid = invalid_indices.shape[0] # p.invalid = np.sum(p.invalid_tokens, dtype=np.float32) # overwrite "True" with the number of invalid tokens return invalid_reward ### Observation noise # For reward_noise_type == "y_hat", success must always be checked to # ensure success cases aren't overlooked due to noise. If successful, # return max_reward. if reward_noise and reward_noise_type == "y_hat": if p.evaluate.get("success"): return max_reward y_hat += rng.normal(loc=0, scale=scale, size=y_hat.shape) # Compute metric # r = 1 / (1 + np.sqrt(np.mean((y_train - y_hat) ** 2) / np.var(y_train))) # note, y_train always be the first entry otherwies the variance of the wrong set is used r = metric(y_train, y_hat) ### Direct reward noise # For reward_noise_type == "r", success can for ~max_reward metrics be # confirmed before adding noise. If successful, must return np.inf to # avoid overlooking success cases. if reward_noise and reward_noise_type == "r": if r >= max_reward - 1e-5 and p.evaluate.get("success"): return np.inf r += rng.normal(loc=0, scale=scale) if normalize_variance: r /= np.sqrt(1 + 12*scale**2) return r def set_ad_traversal(p): # create ad tokens: ad_traversal = create_ad_tokens(p.traversal) # update traversal with error metric p.ad_traversal = ad_metric_traversal[:-1] + ad_traversal + [ad_metric_traversal[-1]] # add ad_const_pos for ad traversal. p.ad_const_pos = [pos + len(ad_metric_traversal[:-1]) for pos in p.const_pos] def reverse_ad(p): # on each evaluation reset the adjoint tokens to 0, except the first to 1. # Compute estimated values (base_r, jac), invalid_indices = p.ad_reverse(X_train) # For invalid expressions, return invalid_reward if p.invalid: if invalid_indices is None: p.invalid_tokens = invalid_indices # if the program is invalid but invalid indices is None the tokens are not logged, set invalid to 1 p.invalid = 1 else: # if invalid weight = 0, p.invalid_tokens will be a dummy array of len == 2 # Adjust invalid_indices from AD_traversal invalid_indices -= ad_metric_start_idx invalid_indices = invalid_indices[invalid_indices >= 0] p.invalid = invalid_indices.shape[0] p.invalid_tokens = invalid_indices # p.invalid_tokens = p.invalid_tokens[ad_metric_start_idx:-1] # p.invalid = np.sum(p.invalid_tokens, dtype=np.float32) # p.invalid = np.sum(p.invalid_tokens[ad_metric_start_idx:-1], dtype=np.float32) return invalid_reward, np.zeros(jac.shape) # set self.r and self.jac at end of this function return base_r, jac def evaluate(p): # Compute predictions on test data y_hat, invalid_indices = p.execute(X_test) if invalid_indices is not None: nmse_test = None nmse_test_noiseless = None success = False else: # NMSE on test data (used to report final error) nmse_test = np.mean((y_test - y_hat)**2) / var_y_test # NMSE on noiseless test data (used to determine recovery) nmse_test_noiseless = np.mean((y_test_noiseless - y_hat)**2) / var_y_test_noiseless # Success is defined by NMSE on noiseless test data below a threshold success = nmse_test_noiseless < threshold info = { "nmse_test" : nmse_test, "nmse_test_noiseless" : nmse_test_noiseless, "success" : success } if extra_metric_test is not None: if p.invalid: m_test = None m_test_noiseless = None else: m_test = metric_test(y_test, y_hat) m_test_noiseless = metric_test(y_test_noiseless, y_hat) info.update( { extra_metric_test : m_test, extra_metric_test + '_noiseless' : m_test_noiseless } ) return info tokens = create_tokens(n_input_var=X_train.shape[1], function_set=function_set, protected=protected) library = Library(tokens) # create secondary library without the tensors to pass to the controllers tens = [] if enforce_sum: tens_idx = [] for idx, val in enumerate(dataset_info['input']): if val in ['T1', 'T2', 'T3', 'T4', 'T5', 'T6', 'T7', 'T8', 'T9', 'T10']: tens.append(val) tens_idx.append(idx) sec_tokens = create_tokens(n_input_var=X_train.shape[1]-len(tens), function_set=function_set, protected=protected) sec_library = Library(sec_tokens) # # add to function set # ad_function_set = function_set + ['n2', 'sum'] # # create library for automatic differentiation in reverse mode # ad_tokens = create_tokens(n_input_var=X_train.shape[1], # function_set=ad_function_set, # protected=protected) # ad_library = Library(ad_tokens) stochastic = reward_noise > 0.0 extra_info = {} task = dsr.task.Task(reward_function=reward, rotate_batch=rotate_batch, data_shuffle=data_shuffle, ad_reverse=reverse_ad, set_ad_traversal=set_ad_traversal, evaluate=evaluate, library=library, sec_library=sec_library, stochastic=stochastic, extra_info=extra_info) return task
data_i['omega_frozen'], normalize=False) nut = data_i['nut_frozen'] k = data_i['k'] bdelta = np.moveaxis(data_i['bDelta'], -1, 0) bij_LEVM = -(nut / k) * Sij bij_LEVM = np.moveaxis(bij_LEVM, -1, 0) bij_corr = bij_LEVM + bdelta trispace_target, cmap_space_target = calc_lumley_cspace(bij_corr, data_i) #first calculate reward on dataset used fro training (without walls) X, y = load_frozen_RANS_dataset(config['task']['dataset']) if isinstance(results, str): dsr_bDelta = eval_string(results, X) # evaluate string else: dsr_bDelta, _ = results['program'].execute(X) if np.isnan(dsr_bDelta).any(): print( f'Best expression of dsr_{results["name"]}_{results["seed"]} returned NaN on {case}' ) return inv_nrmse = 1 / (1 + np.sqrt(np.mean((y - dsr_bDelta)**2)) / np.std(y))
def scatter_results_scalar(results, config): plot_sparta = True X, y = config['task']['dataset'] logdir = config['training']['logdir'] if plot_sparta: inputs = config['task']['dataset_info']['input'] # find grad_u_T1 if 'grad_u_T1' in inputs: grad_u_T1 = X[:, inputs.index('grad_u_T1')] else: dummy_config = config['task'] dummy_config['dataset']['input'] = ['grad_u_T1'] dummy_config['dataset']['skip_wall'] = False grad_u_T1, _ = load_frozen_RANS_dataset(dummy_config['dataset']) # find K if 'k' in inputs: k = X[:, inputs.index('k')] else: dummy_config = config['task'] dummy_config['dataset']['input'] = ['k'] dummy_config['dataset']['skip_wall'] = False k, _ = load_frozen_RANS_dataset(dummy_config['dataset']) Rsparta = 2 * k * grad_u_T1 * 0.93 # sparta_inrmse = 1 / (1 + np.sqrt(np.mean((y-Rsparta)**2))/np.std(y)) # print(sparta_inrmse) # yhat, _ = results['program'].execute(X) reward = results['r'] expression = results['expression'] name = results['name'] seed = results['seed'] filename = f'dsr_{name}_{seed}' fig = plt.figure(figsize=(15, 15), dpi=100) ax = plt.gca() ax.set_xscale('log') ax.set_yscale('log') plt.scatter(y, yhat, s=2) if plot_sparta: plt.scatter(y, Rsparta, s=2, zorder=-1) plt.legend(['DSR', 'sparta']) plt.xlabel('Target (ground truth)') plt.ylabel('DSR model result') plt.title(f'{config["task"]["metric"]} = {reward} \n expression = ' + expression) plt.grid('both') plt.xlim([10e-6, 1]) plt.ylim([10e-6, 1]) plt.savefig(logdir + '/' + filename)
def contourplot_results_tensor(results, config): # re-read config file in output directory to find in and outputs logdir = config['training']['logdir'] with open(logdir + '/config.json', encoding='utf-8') as f: config = json.load(f) name = results['name'] seed = results['seed'] cases = ['PH10595', 'CD12600', 'CBFS13700'] for case in cases: config['task']['dataset']['name'] = case config['task']['dataset']['skip_wall'] = False X, y = load_frozen_RANS_dataset(config['task']['dataset']) yhat, _ = results['program'].execute(X) y = de_flatten_tensor(y) yhat = de_flatten_tensor(yhat) # load mesh data frozen = pickle.load( open(f'turbulence/frozen_data/{case}_frozen_var.p', 'rb')) data_i = frozen['data_i'] mesh_x = data_i['meshRANS'][0, :, :] mesh_y = data_i['meshRANS'][1, :, :] filename = f'{logdir}/dsr_{name}_{seed}_contour_{case}' # here set up grid spec, dress 6 axes one by one using dress_contour_axes fig = plt.figure(figsize=(22, 15), dpi=250) # , constrained_layout=False outer_grid = fig.add_gridspec( 2, 3) # outer_grid = fig11.add_gridspec(4, 4, wspace=0, hspace=0) axes = [] for row in range(2): for col in range(3): inner_grid = outer_grid[row, col].subgridspec(2, 1, wspace=0, hspace=0) axs = inner_grid.subplots() r = row c = col if r == 1: c = col + 1 if (r == 1) and (c == 3): r = 2 c = 2 # outer_grid[row, col].suptitle(f'T{r + 1},{c + 1}') ymin = np.min(y[r, c, :]) ymax = np.max(y[r, c, :]) y_plot = y[r, c, :] yhat_plot = yhat[r, c, :] y_plot = np.reshape(y_plot, mesh_x.shape, order='F') yhat_plot = np.reshape(yhat_plot, mesh_x.shape, order='F') ax0 = axs[0].contourf(mesh_x, mesh_y, y_plot, levels=30, vmin=ymin, vmax=ymax, cmap='Reds') axs[0].set_title(f'Target {r+1},{c+1}', y=1.0, pad=-14) ax1 = axs[1].contourf(mesh_x, mesh_y, yhat_plot, levels=30, vmin=ymin, vmax=ymax, cmap='Reds') axs[1].set_title('DSR Model', y=1.0, pad=-14) axes.append([ax0, ax1, axs]) for pair in axes: ax0 = pair[0] axs = pair[2] fig.colorbar(ax0, ax=axs[0]) fig.colorbar(ax0, ax=axs[1]) plt.savefig(filename, bbox_inches='tight')