gmm.py

# -*- coding: utf-8 -*-
"""
Created on Sun May 24 10:07:00 2020

@author: modified from https://www.katnoria.com/mdn/ , a tutorial on tf2 gdns

Made to fit the Einasto profile data generated by Martin Sanner

Recent changes:
    taken out premultiplier in model

List of outstanding issues:
    
Notes:
    Talk about possibility of learning underlying distribution of inputs, and not the actual function
    
"""

import numpy as np
import matplotlib.pyplot as plt
import EinastoSim
import h5py

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 
import tensorflow as tf
#import tensorflow_addons as tfa #AdamW
import tensorflow_probability as tfp#normal dist
from copy import deepcopy
import sys
import  logging
from datetime import datetime

import pandas as pd

from normal_dist_calculator import generate_tensor_mixture_model
from Reparameterizer import reparameterizer, normalize_profiles,renormalize_profiles

import pickle

import argparse

import trainingAddons as trad


np.random.seed(42)
tf.random.set_seed(42)
plt.close("all")
'''
Logging: Taken from https://stackoverflow.com/a/13733863
'''

now = datetime.now()
d_string = now.strftime("%d/%m/%Y, %H:%M:%S")

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler("logfile_{}_{}.log".format(now.day,now.month)),
        logging.StreamHandler(sys.stdout)
    ]
)
'''

'''

    
if __name__ == "__main__":
    run_file = "./runID.txt"
    run_id = -1
    if not os.path.isfile(run_file):
        with open(run_file,"w") as f:
            run_id = 1
            f.write(str(run_id))
            
    else:
        with open(run_file,"r") as f:
            run_id = int(f.read())
    logging.info("="*20)
    logging.info("Starting new run #{} at {}".format(run_id,d_string))
    logging.info("="*20)
    def restricted_float(x):
        try:
            x = float(x)
        except ValueError:
            raise argparse.ArgumentTypeError("%r not a floating-point literal" % (x,))
        return x
    
    def_num_profiles = 2500
    def_train_ratio = 0.5
    def_lr  = 1e-3
    def_k = 8
    #def_kg = 1
    def_epochs = 1000
    parser = argparse.ArgumentParser()
    parser.add_argument("--num_profile",type = int,default = def_num_profiles, help = "Number of profiles - default {}".format(def_num_profiles))
    parser.add_argument("--train_ratio",type = restricted_float, default = def_train_ratio, help = "Ratio of training to test samples - default {}".format(def_train_ratio))
    parser.add_argument("--lr",type = restricted_float,default = def_lr,help = "Learning rate - default {}".format(def_lr))
    parser.add_argument("--k",type = int, default = def_k, help = "k - default {}".format(def_k))
    #parser.add_argument("--kg",type = int, default = def_kg, help = "k-generator - default {}".format(def_kg))
    parser.add_argument("--epochs",type = int, default = def_epochs, help = "Epochs - default {}".format(def_epochs))
    args = parser.parse_args()
    
    num_profile_train = int(args.num_profile*args.train_ratio)
    

    logging.info("Running on GPU: {}".format(len(tf.config.experimental.list_physical_devices('GPU')) > 0))
    logging.info("Generating {} Profiles for training".format(num_profile_train))
    
    sample_profiles,profile_params,associated_r = EinastoSim.generate_n_random_einasto_profile_maggie(num_profile_train)
    
    logging.info("Defining backend type: TF.float64")
    tf.keras.backend.set_floatx("float64")
    logging.info("Running Logged renormalized Profiles.")
    
    # remove log as test
    sample_profiles_logged = np.asarray([np.log(p) for p in sample_profiles]).astype(np.float64)
    sample_reparam = reparameterizer(sample_profiles_logged)
    sample_profiles_renormed = normalize_profiles(sample_profiles_logged).astype(np.float64)#np.asarray(calculate_renorm_profiles(sample_profiles_logged)).astype(np.float64)
    
    X_full = profile_params
    X_full = np.asarray(X_full).astype(np.float64)
    l = len(profile_params[0])
    
    #output dimension
    out_dim = 1 #just r
    # Number of gaussians to represent the multimodal distribution
    k = args.k
    logging.info("Running {} dimensions on {} distributions".format(out_dim,k))
    
    n_hid_1 = 20
    
    n_hid_2 = 20#int(np.floor((num_profile_train-((k*out_dim*3+n_hid_1*(l+1))))/(n_hid_1+k*out_dim*3+1)))
    logging.info("Hidden Layer sizes: {},{}".format(n_hid_1,n_hid_2))
    
    initial_nodes,best_nodes = trad.create_initial_nodes(l,n_hid_1,n_hid_2,k,out_dim)
    
    
    losses = []
    EPOCHS = args.epochs
    
    lr = args.lr
    optimizer = tf.optimizers.Adam(lr)#tf.optimizers.Adadelta(lr)#tfa.optimizers.AdamW(lr,wd)
    logging.info("Training with optimizer: {}".format(optimizer.__class__.__name__))
    
    
    N = np.asarray(X_full).shape[0]
    num_batches = 10
    
    batchsize = N//num_batches
    logging.info("Employing {} batches with size {}".format(num_batches,batchsize))
    dataset = tf.data.Dataset \
    .from_tensor_slices((X_full, sample_profiles_renormed)) \
    .batch(batchsize)
    
    # Start training
    
    
    n_test_profiles = int(args.num_profile*(1-args.train_ratio))
    train_testing_profile, tt_p_para,t_a_r = EinastoSim.generate_n_random_einasto_profile_maggie(n_test_profiles)
    ttp_logged = np.asarray([np.log(p) for p in train_testing_profile]).astype(np.float64)
    ttp_reparam = reparameterizer(ttp_logged)
    ttp_renormed = normalize_profiles(ttp_logged).astype(np.float64)
    X_tt = tt_p_para
    
    counter_max = 500
    
    loss_target = 1e-3
    weight_decay = 1e-3
    min_epoch_train_pre_div = 300
    '''
    Start_parameters dict keys:
    Counter_max
    loss_target
    max_diff
    epoch
    test_MAEs
    counters
    MSEs
    minDelta
    losses
    '''
    start_parameters = {}
    normalize = False
    max_loss_divergence = 10
    patience_disabled = False
    print_every_n_epochs = 20
    
    best_nodes,losses,MSEs,counters,test_MAEs = trad.train_model(initial_nodes,
                                                                 optimizer,
                                                                 dataset,
                                                                 associated_r,
                                                                 EPOCHS,
                                                                 X_tt,
                                                                 ttp_renormed,
                                                                 t_a_r,
                                                                 min_epoch_train_pre_div,
                                                                 start_parameters,
                                                                 print_every_n_epochs,
                                                                 lr,
                                                                 normalize,
                                                                 max_loss_divergence,
                                                                 patience_disabled,
                                                                 weight_decay)
    
    
    
    plot_folder = ".//plots//Run_{}//".format(run_id)
    save_folder = ".//models//Run_{}//best_model".format(run_id)
    data_folder = ".//data//Run_{}//".format(run_id)
    if not os.path.exists(data_folder):
        os.makedirs(data_folder)
        
    #logging.info("Saving best model to {}".format(save_folder))
    #best_model.save_weights(save_folder)
                
    logging.info("Dumping data to {}".format(data_folder))
    now = datetime.now()
    with open(data_folder+"MAE_Losses.dat","wb") as f:
        pickle.dump(losses,f)
    with open(data_folder+"MSE_Losses.dat","wb") as f:
        pickle.dump(MSEs,f)
    with open(data_folder+"Patience.dat","wb") as f:
        pickle.dump(counters,f)
    with open(data_folder+"mae_test_losses.dat","wb") as f:
        pickle.dump(test_MAEs,f)
    
    n_test_profiles = 100
    test_profiles,t_profile_params,t_associated_r = EinastoSim.generate_n_random_einasto_profile_maggie(n_test_profiles)
    t_sample_profiles_logged = np.asarray([np.log(p) for p in test_profiles]).astype(np.float64)
    t_s_reparam = reparameterizer(t_sample_profiles_logged)
    t_s_renorm = normalize_profiles(t_sample_profiles_logged).astype(np.float64)
    X_test = t_profile_params
    
    pi_test, mu_test,var_test = trad.model(np.asarray(X_test),best_nodes,normalize)
    #var_test = tf.exp(var_test_log)
    test_data = {"Profiles":t_s_renorm, "STDParams":{"Pi":pi_test,"Mu":mu_test,"Var":var_test},"Xtest":X_test, "r":t_associated_r}
    with open(data_folder+"test_data.dat","wb") as f:
        pickle.dump(test_data,f)
    
    with open(run_file,"w") as f:
        f.write(str(run_id +1))