Module tools

# -*- encoding=utf-8 -*-

from pathlib import Path
import shutil, os, pickle, math
import numpy as np
import matplotlib.pyplot as plt

# own
from pf_to_dem import *

#------------------------------------------------------------------------------------------------------------------------------------------ #

def create_folder(name):
    '''
    Create a new folder. If it already exists, it is erased.
    '''
    if Path(name).exists():
        shutil.rmtree(name)
    os.mkdir(name)

#------------------------------------------------------------------------------------------------------------------------------------------ #

def save_mesh_database(dict_user, dict_sample):
    '''
    Save mesh database.
    '''
    # creating a database
    if not Path('mesh_map.database').exists():
        dict_data = {
        'n_proc': dict_user['n_proc'],
        'n_mesh_x': len(dict_sample['x_L']),
        'n_mesh_y': len(dict_sample['y_L']),
        'n_mesh_z': len(dict_sample['z_L']),
        'L_L_i_XYZ_used': dict_sample['L_L_i_XYZ_used'],
        }
        dict_database = {'Run_1': dict_data}
        with open('mesh_map.database', 'wb') as handle:
                pickle.dump(dict_database, handle, protocol=pickle.HIGHEST_PROTOCOL)
    # updating a database
    else :
        with open('mesh_map.database', 'rb') as handle:
            dict_database = pickle.load(handle)
        dict_data = {
        'n_proc': dict_user['n_proc'],
        'n_mesh_x': len(dict_sample['x_L']),
        'n_mesh_y': len(dict_sample['y_L']),
        'n_mesh_z': len(dict_sample['z_L']),
        'L_L_i_XYZ_used': dict_sample['L_L_i_XYZ_used']
        }   
        mesh_map_known = False
        for i_run in range(1,len(dict_database.keys())+1):
            if dict_database['Run_'+str(int(i_run))]['n_proc'] == dict_data['n_proc'] and \
               dict_database['Run_'+str(int(i_run))]['n_mesh_x'] == dict_data['n_mesh_x'] and \
               dict_database['Run_'+str(int(i_run))]['n_mesh_y'] == dict_data['n_mesh_y'] and \
               dict_database['Run_'+str(int(i_run))]['n_mesh_z'] == dict_data['n_mesh_z']:
                mesh_map_known = True
        # new entry
        if not mesh_map_known: 
            key_entry = 'Run_'+str(int(len(dict_database.keys())+1))
            dict_database[key_entry] = dict_data
            with open('mesh_map.database', 'wb') as handle:
                pickle.dump(dict_database, handle, protocol=pickle.HIGHEST_PROTOCOL)

#------------------------------------------------------------------------------------------------------------------------------------------ #

def check_mesh_database(dict_user, dict_sample):
    '''
    Check mesh database.
    '''
    if Path('mesh_map.database').exists():
        with open('mesh_map.database', 'rb') as handle:
            dict_database = pickle.load(handle)
        dict_data = {
        'n_proc': dict_user['n_proc'],
        'n_mesh_x': len(dict_sample['x_L']),
        'n_mesh_y': len(dict_sample['y_L']),
        'n_mesh_z': len(dict_sample['z_L'])
        }   
        mesh_map_known = False
        for i_run in range(1,len(dict_database.keys())+1):
            if dict_database['Run_'+str(int(i_run))]['n_proc'] == dict_data['n_proc'] and \
               dict_database['Run_'+str(int(i_run))]['n_mesh_x'] == dict_data['n_mesh_x'] and \
               dict_database['Run_'+str(int(i_run))]['n_mesh_y'] == dict_data['n_mesh_y'] and \
               dict_database['Run_'+str(int(i_run))]['n_mesh_z'] == dict_data['n_mesh_z']:
                mesh_map_known = True
                i_known = i_run
        if mesh_map_known :
            dict_sample['Map_known'] = True
            dict_sample['L_L_i_XYZ_used'] = dict_database['Run_'+str(int(i_known))]['L_L_i_XYZ_used']
        else :
            dict_sample['Map_known'] = False
    else :
        dict_sample['Map_known'] = False

#------------------------------------------------------------------------------------------------------------------------------------------ #

def plot_sum_mean_etai_c(dict_user, dict_sample):
    '''
    Plot figure illustrating the sum and the mean of etai and c.
    '''
    # compute tracker
    # sum
    s_eta_i = 0
    if dict_user['L_L_sum_eta_i'] == []:
        for i_grain in range(len(dict_sample['L_etai_map'])):
            dict_user['L_L_sum_eta_i'].append([np.sum(dict_sample['L_etai_map'][i_grain])])
            s_eta_i = s_eta_i + np.sum(dict_sample['L_etai_map'][i_grain])
    else :
        for i_grain in range(len(dict_sample['L_etai_map'])):
            dict_user['L_L_sum_eta_i'][i_grain].append(np.sum(dict_sample['L_etai_map'][i_grain]))
            s_eta_i = s_eta_i + np.sum(dict_sample['L_etai_map'][i_grain])
    dict_user['L_sum_c'].append(np.sum(dict_sample['c_map']))
    dict_user['L_sum_mass'].append(1/dict_user['V_m']*s_eta_i+np.sum(dict_sample['c_map']))
    # mean
    m_eta_i = 0
    if dict_user['L_L_m_eta_i'] == []:
        for i_grain in range(len(dict_sample['L_etai_map'])):
            dict_user['L_L_m_eta_i'].append([np.mean(dict_sample['L_etai_map'][i_grain])])
            m_eta_i = m_eta_i + np.mean(dict_sample['L_etai_map'][i_grain])
    else :
        for i_grain in range(len(dict_sample['L_etai_map'])):
            dict_user['L_L_m_eta_i'][i_grain].append(np.mean(dict_sample['L_etai_map'][i_grain]))
            m_eta_i = m_eta_i + np.mean(dict_sample['L_etai_map'][i_grain])
    dict_user['L_m_c'].append(np.mean(dict_sample['c_map']))
    dict_user['L_m_mass'].append(1/dict_user['V_m']*m_eta_i+np.mean(dict_sample['c_map']))

    # plot sum eta_i, c
    if 'sum_etai_c' in dict_user['L_figures']:
        fig, ((ax1,ax2),(ax3,ax4)) = plt.subplots(nrows=2,ncols=2,figsize=(16,9))
        for i_grain in range(len(dict_sample['L_L_sum_eta_i'])):
            ax1.plot(dict_user['L_L_sum_eta_i'][i_grain])
        ax1.set_title(r'$\Sigma\eta_i$')
        ax3.plot(dict_user['L_sum_c'])
        ax3.set_title(r'$\Sigma C$')
        ax4.plot(dict_user['L_sum_mass'])
        ax4.set_title(r'$\Sigma mass$')
        fig.tight_layout()
        fig.savefig('plot/sum_etai_c.png')
        plt.close(fig)

    # plot mean eta_i, c
    if 'mean_etai_c' in dict_user['L_figures']:
        fig, ((ax1,ax2),(ax3,ax4)) = plt.subplots(nrows=2,ncols=2,figsize=(16,9))
        for i_grain in range(len(dict_user['L_L_m_eta_i'])):
            ax1.plot(dict_user['L_L_m_eta_i'][i_grain])
        ax1.set_title(r'Mean $\eta_i$')
        ax3.plot(dict_user['L_m_c'])
        ax3.set_title(r'Mean $c$')
        ax4.plot(dict_user['L_m_mass'])
        ax4.set_title(r'Mean mass')
        fig.tight_layout()
        fig.savefig('plot/mean_etai_c.png')
        plt.close(fig)

#------------------------------------------------------------------------------------------------------------------------------------------ #

def compute_mass(dict_user, dict_sample):
    '''
    Compute the mass at a certain time.
     
    Mass is sum of etai and c.
    '''
    # sum of masses
    L_sum_eta_i_tempo = []
    s_eta_i = 0
    for i_grain in range(len(dict_sample['L_etai_map'])):
        L_sum_eta_i_tempo.append([np.sum(dict_sample['L_etai_map'][i_grain])])
        s_eta_i = s_eta_i + np.sum(dict_sample['L_etai_map'][i_grain])
    dict_user['L_sum_eta_i_tempo'] = L_sum_eta_i_tempo
    dict_user['sum_c_tempo'] = np.sum(dict_sample['c_map'])
    dict_user['sum_mass_tempo'] = 1/dict_user['V_m']*s_eta_i+np.sum(dict_sample['c_map'])

#------------------------------------------------------------------------------------------------------------------------------------------ #

def compute_mass_loss(dict_user, dict_sample, tracker_key):
    '''
    Compute the mass loss from the previous compute_mass() call.
     
    Plot in the given tracker.
    Mass is sum of etai and c.
    '''
    # delta masses
    if dict_user['L_'+tracker_key+'_eta_i'] == []:
        s_eta_i = 0
        for i_grain in range(len(dict_sample['L_etai_map'])):
            detai = np.sum(dict_sample['L_etai_map'][i_grain]) - dict_user['L_sum_eta_i_tempo'][i_grain]
            s_eta_i = s_eta_i + np.sum(dict_sample['L_etai_map'][i_grain])
            # save
            dict_user['L_'+tracker_key+'_eta_i'].append([detai])
    else : 
        s_eta_i = 0
        for i_grain in range(len(dict_sample['L_etai_map'])):
            detai = np.sum(dict_sample['L_etai_map'][i_grain]) - dict_user['L_sum_eta_i_tempo'][i_grain]
            s_eta_i = s_eta_i + np.sum(dict_sample['L_etai_map'][i_grain])
            # save
            dict_user['L_'+tracker_key+'_eta_i'][i_grain].append(detai)
    dc = np.sum(dict_sample['c_map']) - dict_user['sum_c_tempo']
    dm = 1/dict_user['V_m']*s_eta_i+np.sum(dict_sample['c_map']) - dict_user['sum_mass_tempo']
    # save
    dict_user[tracker_key+'_c'].append(dc)
    dict_user[tracker_key+'_m'].append(dm)

    # plot
    if 'mass_loss' in dict_user['L_figures']:
        fig, ((ax1,ax2),(ax3,ax4)) = plt.subplots(nrows=2,ncols=2,figsize=(16,9))
        for i_grain in range(len(dict_sample['L_'+tracker_key+'_eta_i'])):
            ax1.plot(dict_user['L_'+tracker_key+'_eta_i'][i_grain])
        ax1.set_title(r'$\eta_i$ loss')
        ax3.plot(dict_user[tracker_key+'_c'])
        ax3.set_title(r'$c$ loss')
        ax4.plot(dict_user[tracker_key+'_m'])
        ax4.set_title(r'mass loss')
        fig.tight_layout()
        fig.savefig('plot/'+tracker_key+'.png')
        plt.close(fig)

#------------------------------------------------------------------------------------------------------------------------------------------ #

def plot_performances(dict_user, dict_sample):
    '''
    Plot figure illustrating the time performances of the algorithm.
    '''
    if 'performances' in dict_user['L_figures']:
        fig, (ax1) = plt.subplots(nrows=1,ncols=1,figsize=(16,9))
        ax1.plot(dict_user['L_t_dem'], label='DEM')
        ax1.plot(dict_user['L_t_pf'], label='PF')
        ax1.plot(dict_user['L_t_dem_to_pf'], label='DEM to PF')
        ax1.plot(dict_user['L_t_pf_to_dem_1'], label='PF to DEM 1')
        ax1.plot(dict_user['L_t_pf_to_dem_2'], label='PF to DEM 2')
        ax1.legend()
        ax1.set_title('Performances (s)')
        ax1.set_xlabel('Iterations (-)')
        fig.tight_layout()
        fig.savefig('plot/performances.png')
        plt.close(fig)
                
#------------------------------------------------------------------------------------------------------------------------------------------ #

def plot_dem(dict_user, dict_sample):
    '''
    Plot figure illustrating the overlap and force transmitted.
    '''
    # Need to be adapted to multiple contacts  

    if 'overlap' in dict_user['L_figures']:
        fig, (ax1) = plt.subplots(nrows=1,ncols=1,figsize=(16,9))
        for i_contact in range(len(dict_user['L_L_overlap'])):
            ax1.plot(dict_user['L_L_overlap'][i_contact])
        ax1.set_xlabel('iterations (-)')
        ax1.set_ylabel('overlap in DEM (-)')
        fig.tight_layout()
        fig.savefig('plot/dem_overlap.png')
        plt.close(fig)
    if 'normal_force' in dict_user['L_figures']:
        fig, (ax1) = plt.subplots(nrows=1,ncols=1,figsize=(16,9))
        for i_contact in range(len(dict_user['L_L_normal_force'])):
            ax1.plot(dict_user['L_L_normal_force'][i_contact])
        ax1.plot([0, len(dict_user['L_normal_force'])-1],\
                 [dict_user['force_applied'], dict_user['force_applied']], color='k', label='target')
        ax1.set_xlabel('iterations (-)')
        ax1.set_ylabel('normal force (-)')
        fig.tight_layout()
        fig.savefig('plot/dem_normal_force.png')
        plt.close(fig)

#------------------------------------------------------------------------------------------------------------------------------------------ #

def plot_contact(dict_user, dict_sample):
    '''
    Plot figure illustrating the contact characteristics.
    '''
    # Need to be adapted to multiple contacts

    if 'contact_box' in dict_user['L_figures']:
        fig, (ax1,ax2,ax3) = plt.subplots(nrows=1,ncols=3,figsize=(16,9))
        for i_contact in range(len(dict_user['L_L_contact_box_x'])):
            ax1.plot(dict_user['L_L_contact_box_x'][i_contact])
            ax2.plot(dict_user['L_L_contact_box_y'][i_contact])
            ax3.plot(dict_user['L_L_contact_box_z'][i_contact])
        ax1.set_suptitle('x')
        ax2.set_suptitle('y')
        ax3.set_suptitle('z')
        ax1.set_ylabel('contact box dimensions (-)')
        ax2.set_xlabel('iterations (-)')
        fig.tight_layout()
        fig.savefig('plot/contact_box.png')
        plt.close(fig)
    if 'contact_volume' in dict_user['L_figures']:
        fig, (ax1) = plt.subplots(nrows=1,ncols=1,figsize=(16,9))
        for i_contact in range(len(dict_user['L_L_contact_volume'])):
            ax1.plot(dict_user['L_L_contact_volume'][i_contact])
        ax1.set_xlabel('iterations (-)')
        ax1.set_ylabel('contact volume (-)')
        fig.tight_layout()
        fig.savefig('plot/contact_volume.png')
        plt.close(fig)
    if 'contact_surface' in dict_user['L_figures']:
        fig, (ax1) = plt.subplots(nrows=1,ncols=1,figsize=(16,9))
        for i_contact in range(len(dict_user['L_L_contact_surface'])):
            ax1.plot(dict_user['L_L_contact_surface'][i_contact])
        ax1.set_xlabel('iterations (-)')
        ax1.set_ylabel('contact surface (-)')
        fig.tight_layout()
        fig.savefig('plot/contact_surface.png')
        plt.close(fig)

#------------------------------------------------------------------------------------------------------------------------------------------ #

def plot_as_pressure(dict_user, dict_sample):
    '''
    Plot figure illustrating the solid activity and pressure at the contact.
    '''
    if 'as' in dict_user['L_figures']:
        fig, (ax1) = plt.subplots(nrows=1,ncols=1,figsize=(16,9))
        for i_contact in range(len(dict_user['L_L_contact_as'])):
            ax1.plot(dict_user['L_L_contact_as'][i_contact])
        ax1.set_xlabel('iterations (-)')
        ax1.set_ylabel('solid activity (-)')
        fig.tight_layout()
        fig.savefig('plot/contact_as.png')
        plt.close(fig)
    if 'pressure' in dict_user['L_figures']:
        fig, (ax1) = plt.subplots(nrows=1,ncols=1,figsize=(16,9))
        for i_contact in range(len(dict_user['L_L_contact_pressure'])):
            ax1.plot(dict_user['L_L_contact_pressure'][i_contact])
        ax1.set_xlabel('iterations (-)')
        ax1.set_ylabel('solid pressure (-)')
        fig.tight_layout()
        fig.savefig('plot/contact_pressure.png')
        plt.close(fig)

#------------------------------------------------------------------------------------------------------------------------------------------ #

def plot_slices(dict_user, dict_sample):
    '''
    Plot slices of the configuration.
    '''
    # determine the index of the slices
    i_x_slice = int(len(dict_sample['x_L'])/2)
    i_y_slice = int(len(dict_sample['y_L'])/2)
    i_z_slice = int(len(dict_sample['z_L'])/2)

    if 'configuration_eta' in dict_user['L_figures'] :
        for i_grain in range(len(dict_sample['L_etai_map'])):
            # plot
            fig, (ax1, ax2, ax3) = plt.subplots(nrows=1,ncols=3,figsize=(16,9))
            # y-z
            im = ax1.imshow(dict_sample['L_etai_map'][i_grain][i_x_slice,:,:], interpolation = 'nearest')
            fig.colorbar(im, ax=ax1)
            ax1.set_title(r'slice y-z',fontsize = 30)
            # x-z
            im = ax2.imshow(dict_sample['L_etai_map'][i_grain][:,i_y_slice,:], interpolation = 'nearest')
            fig.colorbar(im, ax=ax2)
            ax2.set_title(r'slice x-z',fontsize = 30)
            # x-y
            im = ax3.imshow(dict_sample['L_etai_map'][i_grain][:,:,i_z_slice], interpolation = 'nearest')
            fig.colorbar(im, ax=ax3)
            ax3.set_title(r'slice x-y',fontsize = 30)
            # close
            fig.tight_layout()
            fig.savefig('plot/configuration/eta_'+str(i_grain+1)+'_'+str(dict_sample['i_DEMPF_ite'])+'.png')
            plt.close(fig)

    if 'configuration_c' in dict_user['L_figures']:
        # plot
        fig, (ax1, ax2, ax3) = plt.subplots(nrows=1,ncols=3,figsize=(16,9))
        # y-z
        im = ax1.imshow(dict_sample['c_map'][i_x_slice,:,:], interpolation = 'nearest')
        fig.colorbar(im, ax=ax1)
        ax1.set_title(r'slice y-z',fontsize = 30)
        # x-z
        im = ax2.imshow(dict_sample['c_map'][:,i_y_slice,:], interpolation = 'nearest')
        fig.colorbar(im, ax=ax2)
        ax2.set_title(r'slice x-z',fontsize = 30)
        # x-y
        im = ax3.imshow(dict_sample['c_map'][:,:,i_z_slice], interpolation = 'nearest')
        fig.colorbar(im, ax=ax3)
        ax3.set_title(r'slice x-y',fontsize = 30)
        # close
        fig.tight_layout()
        fig.savefig('plot/configuration/c_'+str(dict_sample['i_DEMPF_ite'])+'.png')
        plt.close(fig)

    if 'configuration_s_eta' in dict_user['L_figures']:
        # initialization
        map_sum_etas = np.zeros((dict_sample['c_map'].shape[0], dict_sample['c_map'].shape[2]))
        # iterate on grains
        for i_grain in range(len(dict_sample['L_etai_map'])):
            # x-z 
            map_sum_etas = map_sum_etas + dict_sample['L_etai_map'][i_grain][:,i_y_slice,:]
        # plot
        fig, (ax1) = plt.subplots(nrows=1,ncols=1,figsize=(16,9))
        im = ax1.imshow(map_sum_etas, interpolation = 'nearest')
        ax1.set_title(r'slice x-z',fontsize = 30)
        fig.colorbar(im, ax=ax1)
        fig.tight_layout()
        fig.savefig('plot/configuration/s_eta_'+str(dict_sample['i_DEMPF_ite'])+'.png')
        plt.close(fig)        

#------------------------------------------------------------------------------------------------------------------------------------------ #

def plot_displacement(dict_user, dict_sample):
    '''
    Plot figure illustrating the cumulative displacement.
    '''
    # pp data
    L_L_strain = []
    for i_grain in range(len(dict_user['L_L_displacement'][0])):
        L_strain = []
        for i_displacement in range(len(dict_user['L_L_displacement'])):
            if i_displacement == 0:
                L_displacement_cum = [dict_user['L_L_displacement'][i_displacement][i_grain][2]]
            else : 
                L_displacement_cum.append(L_displacement_cum[-1]+dict_user['L_L_displacement'][i_displacement][i_grain][2])
            L_strain.append(L_displacement_cum[-1]/(2*dict_user['radius']))
        L_L_strain.append(L_strain)
    # plot
    fig, (ax1) = plt.subplots(nrows=1,ncols=1,figsize=(16,9))
    for i_grain in range(len(L_L_strain)):
        ax1.plot(L_L_strain[i_grain])
    ax1.set_xlabel('iterations (-)')
    ax1.set_ylabel('vertical strain (-)')
    fig.tight_layout()
    fig.savefig('plot/vertical_strain_grain.png')
    plt.close(fig)

#------------------------------------------------------------------------------------------------------------------------------------------ #

def plot_settlement(dict_user, dict_sample):
    '''
    Plot figure illustrating the settlement curve.
    '''
    # pp data
    L_strain = []
    for i_size in range(len(dict_user['L_delta_z_sample'])):
        L_strain.append((dict_user['L_delta_z_sample'][i_size]-dict_user['L_delta_z_sample'][0])/dict_user['L_delta_z_sample'][0]) 
    # plot
    fig, (ax1) = plt.subplots(nrows=1,ncols=1,figsize=(16,9))
    ax1.plot(L_strain)
    ax1.set_xlabel('iterations (-)')
    ax1.set_ylabel('vertical strain (-)')
    fig.tight_layout()
    fig.savefig('plot/vertical_strain.png')
    plt.close(fig)