267 lines
		
	
	
		
			10 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
			
		
		
	
	
			267 lines
		
	
	
		
			10 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
#!/usr/bin/env python
 | 
						|
# -*- coding: utf-8 -*-
 | 
						|
import os
 | 
						|
import glob
 | 
						|
import subprocess
 | 
						|
 | 
						|
import json
 | 
						|
 | 
						|
import numpy as np
 | 
						|
import numpy.polynomial.polynomial as poly
 | 
						|
 | 
						|
from scipy.sparse import spdiags
 | 
						|
from scipy.optimize import curve_fit
 | 
						|
 | 
						|
import matplotlib.pyplot as plt
 | 
						|
 | 
						|
from itertools import cycle
 | 
						|
 | 
						|
from pylot.tomography.fmtomo_tools.fmtomo_grid_utils import read_vgrid
 | 
						|
 | 
						|
# def calc_dampnorm(vgrid, vgrid_ref):
 | 
						|
#     # calculate m - m0
 | 
						|
#     m_m0 = np.array(vgrid['vps']) - np.array(vgrid_ref['vps'])
 | 
						|
#
 | 
						|
#     # calculate inverse of diagonal elements of a priori model covariance matrix (which should be a diagonal matrix)
 | 
						|
#     # IMPORTANT: COVARIANCES ARE MOST LIKELY STANDARD DEVIATIONS -> square them
 | 
						|
#     covs = np.array(vgrid_ref['covs'])**2
 | 
						|
#
 | 
						|
#     covs_inv = 1. / covs
 | 
						|
#
 | 
						|
#     #cm_inv = spdiags(covs_inv, 0, covs_inv.size, covs_inv.size)#
 | 
						|
#
 | 
						|
#     #norm_calc_old = np.dot(m_m0.transpose(), m_m0 * cm_inv)
 | 
						|
#
 | 
						|
#     norm = np.dot(m_m0, m_m0 * covs_inv)
 | 
						|
#
 | 
						|
#     return norm
 | 
						|
#
 | 
						|
#
 | 
						|
# def calc_smoothnorm(vgrid, gridn, R=6371.):
 | 
						|
#     nR, nTheta, nPhi = gridn
 | 
						|
#
 | 
						|
#     vps = np.array(vgrid['vps'])
 | 
						|
#     lats = np.array(vgrid['lats'])
 | 
						|
#     lons = np.array(vgrid['lons'])
 | 
						|
#     depths = np.array(vgrid['depths'])
 | 
						|
#     #vgridref = np.array(vgridref['vps'])
 | 
						|
#     #dvgrid = vgrid - vgridref
 | 
						|
#
 | 
						|
#     vgarray = np.zeros((nR, nTheta, nPhi))
 | 
						|
#     lonsarray_km = np.zeros((nR, nTheta, nPhi))
 | 
						|
#     latsarray_km = np.zeros((nR, nTheta, nPhi))
 | 
						|
#     radsarray_km = np.zeros((nR, nTheta, nPhi))
 | 
						|
#     #vgarray_diff = np.zeros((nR, nTheta, nPhi))
 | 
						|
#     smootharray = np.zeros((nR, nTheta, nPhi))
 | 
						|
#     #for iLayer in range(nlayers):
 | 
						|
#     globInd = 0
 | 
						|
#     for iR in range(nR):
 | 
						|
#         for iTheta in range(nTheta):
 | 
						|
#             for iPhi in range(nPhi):
 | 
						|
#                 r = R - depths[globInd]
 | 
						|
#                 lat = lats[globInd]
 | 
						|
#                 lon = lons[globInd]
 | 
						|
#                 r_minor = np.cos(np.deg2rad(lat)) * r
 | 
						|
#                 vgarray[iR, iTheta, iPhi] = vps[globInd]
 | 
						|
#                 radsarray_km[iR, iTheta, iPhi] = r
 | 
						|
#                 latsarray_km[iR, iTheta, iPhi] = np.pi * r * lat / 180.
 | 
						|
#                 lonsarray_km[iR, iTheta, iPhi] = np.pi * r_minor * lon / 180.
 | 
						|
#                 #vgarray_diff[iR, iTheta, iPhi] = vgrid[globInd]
 | 
						|
#                 globInd += 1
 | 
						|
#
 | 
						|
#     # iterate over grid diffs (correct?) and sum 1 * point left -2 * point + 1 * point right in all 3 dim.
 | 
						|
#     smsum = 0.
 | 
						|
#     for iR in range(nR):
 | 
						|
#         for iTheta in range(nTheta):
 | 
						|
#             for iPhi in range(nPhi):
 | 
						|
#                 vg = vgarray[iR, iTheta, iPhi]
 | 
						|
#                 sum1 = sum2 = sum3 = 0.
 | 
						|
#                 if 0 < iPhi < nPhi - 1:
 | 
						|
#                     h = abs(lonsarray_km[iR, iTheta, iPhi + 1] - lonsarray_km[iR, iTheta, iPhi - 1]) / 2
 | 
						|
#                     sum1 = (vgarray[iR, iTheta, iPhi - 1] - 2 * vg + vgarray[iR, iTheta, iPhi + 1]) / h**2
 | 
						|
#                 if 0 < iTheta < nTheta - 1:
 | 
						|
#                     h = abs(latsarray_km[iR, iTheta + 1, iPhi] - latsarray_km[iR, iTheta - 1, iPhi]) / 2
 | 
						|
#                     sum2 = (vgarray[iR, iTheta - 1, iPhi] - 2 * vg + vgarray[iR, iTheta + 1, iPhi]) / h**2
 | 
						|
#                 if 0 < iR < nR - 1:
 | 
						|
#                     h = abs(radsarray_km[iR - 1, iTheta, iPhi] - radsarray_km[iR + 1, iTheta, iPhi]) / 2
 | 
						|
#                     sum3 = (vgarray[iR - 1, iTheta, iPhi] - 2 * vg + vgarray[iR + 1, iTheta, iPhi]) / h**2
 | 
						|
#                 smsum += np.sqrt(sum1**2 + sum2**2 + sum3**2)
 | 
						|
#                 #print(sum1, sum2, sum3, smsum)
 | 
						|
#                 smootharray[iR, iTheta, iPhi] = smsum#sum1 + sum2 + sum3
 | 
						|
#
 | 
						|
#     # m_T * D_T * D * m ?? todo: unsure
 | 
						|
#     norm = np.sum(smootharray ** 2)
 | 
						|
#
 | 
						|
#     return norm, smootharray
 | 
						|
from pylot.tomography.utils import normed_figure
 | 
						|
 | 
						|
 | 
						|
def calc_smoothnorm(wdir, iter):
 | 
						|
    smv = np.loadtxt(os.path.join(wdir, 'it_{}/smv.out'.format(iter + 1)), skiprows=1)
 | 
						|
    dm = np.loadtxt(os.path.join(wdir, 'it_{}/dm.out'.format(iter + 1)), skiprows=1)
 | 
						|
    norm = np.sum(smv*dm)
 | 
						|
    return norm
 | 
						|
 | 
						|
 | 
						|
def calc_dampnorm(wdir, iter):
 | 
						|
    ecmi = np.loadtxt(os.path.join(wdir, 'it_{}/ecmi.out'.format(iter + 1)), skiprows=1)
 | 
						|
    dm = np.loadtxt(os.path.join(wdir, 'it_{}/dm.out'.format(iter + 1)), skiprows=1)
 | 
						|
    norm = np.sum(ecmi * dm**2)
 | 
						|
    return norm
 | 
						|
 | 
						|
 | 
						|
def calc_norm(wdir, iteration_number):
 | 
						|
    dampnorm = calc_dampnorm(wdir, iteration_number)
 | 
						|
    smoothnorm = calc_smoothnorm(wdir, iteration_number)
 | 
						|
 | 
						|
    print('dampnorm: ', dampnorm)
 | 
						|
    print('smoothnorm: ', smoothnorm)
 | 
						|
 | 
						|
    norm = dampnorm + smoothnorm
 | 
						|
 | 
						|
    print('Calculated summed norm of', norm)
 | 
						|
 | 
						|
    return norm, dampnorm, smoothnorm
 | 
						|
 | 
						|
 | 
						|
def calc_tradeoff(fpath_in, fname_out=None, iteration_number = 12):
 | 
						|
    results = {}
 | 
						|
 | 
						|
    for wdir in glob.glob(fpath_in):
 | 
						|
        #wdir = '/rscratch/minos13/marcel/fmtomo_alparray/alparray_mantle_from_m6.0_diehl_crustal_corrections_sm1000_damp100/'
 | 
						|
        smooth = float(wdir.split('_')[-2].split('sm')[-1])
 | 
						|
        damp = float(wdir.split('_damp')[-1].split('/')[0])
 | 
						|
 | 
						|
        print('Calculating tradeoff for smoothing and damping of {}, {}'.format(smooth, damp))
 | 
						|
        if not smooth in results.keys():
 | 
						|
            results[smooth] = {}
 | 
						|
 | 
						|
        iteration_number_new = iteration_number
 | 
						|
        ecmi_path = os.path.join(wdir, 'it_{}'.format(iteration_number_new + 1), 'ecmi.out')
 | 
						|
        smv_path = os.path.join(wdir, 'it_{}'.format(iteration_number_new + 1), 'smv.out')
 | 
						|
        while not os.path.isfile(ecmi_path) or not os.path.isfile(smv_path):
 | 
						|
            iteration_number_new -= 1
 | 
						|
            ecmi_path = os.path.join(wdir, 'it_{}'.format(iteration_number_new + 1), 'ecmi.out')
 | 
						|
            smv_path = os.path.join(wdir, 'it_{}'.format(iteration_number_new + 1), 'smv.out')
 | 
						|
            print('WARNING: Iteration number lowered by 1:', iteration_number_new)
 | 
						|
            if iteration_number_new <= 1:
 | 
						|
                break
 | 
						|
 | 
						|
        if iteration_number_new <= 1:
 | 
						|
            continue
 | 
						|
        else:
 | 
						|
            iteration_number = iteration_number_new
 | 
						|
 | 
						|
        #vgrid, gridn, griddelta, gridstart = read_vgrid(vgrid_path)
 | 
						|
        #vgrid_ref = read_vgrid(os.path.join(wdir, 'vgridsref.in'))[0]
 | 
						|
 | 
						|
        norm, dampnorm, smoothnorm = calc_norm(wdir, iteration_number)
 | 
						|
 | 
						|
        try:
 | 
						|
            fpath = os.path.join(wdir, 'residuals.dat')
 | 
						|
            chi = float(subprocess.check_output(['tail', fpath]).split()[-1])
 | 
						|
        except Exception as e:
 | 
						|
            print(e)
 | 
						|
            chi = np.nan
 | 
						|
 | 
						|
        results[smooth][wdir] = {'dampnorm': dampnorm, 'smoothnorm': smoothnorm,
 | 
						|
                                 'norm': norm, 'chi': chi, 'damp': damp}
 | 
						|
 | 
						|
    #print some output
 | 
						|
    for smooth, result in results.items():
 | 
						|
        print('Smoothing:', smooth)
 | 
						|
        for wdir, item in result.items():
 | 
						|
            print(item['chi'], item['norm'])
 | 
						|
            print(20*'#')
 | 
						|
 | 
						|
    if fname_out:
 | 
						|
        with open(fname_out, 'w') as outfile:
 | 
						|
            json.dump(results, outfile)
 | 
						|
    return results
 | 
						|
 | 
						|
def quadratic_function(x, a, b, c):
 | 
						|
    return a * x ** 2 + b * x + c
 | 
						|
 | 
						|
def one_over_x(x, a, b, c):
 | 
						|
    return a / (x - b) + c
 | 
						|
 | 
						|
def exp_func(x, a, b, c):
 | 
						|
    return a * np.exp(-b * x) + c
 | 
						|
 | 
						|
def plot_tradeoff(fname_in, fix='smooth', plot_norm='both', min_smooth=0, min_damp=0, max_smooth=1e6, max_damp=1e6):
 | 
						|
    with open(fname_in, 'r') as infile:
 | 
						|
        results_smooth = json.load(infile)
 | 
						|
 | 
						|
    lines = ["-", "--", "-.", ":"]
 | 
						|
    linecycler = cycle(lines)
 | 
						|
 | 
						|
    # array will be built for each line: (smooth, damp, norm, chi)
 | 
						|
    plot_values = []
 | 
						|
    for smooth, result in results_smooth.items():
 | 
						|
        for item in result.values():
 | 
						|
            smooth = float(smooth)
 | 
						|
            damping = item['damp']
 | 
						|
            if smooth < min_smooth or damping < min_damp or smooth > max_smooth or damping > max_damp:
 | 
						|
                continue
 | 
						|
            plot_values.append(np.array([smooth, damping, item[plot_norm], item['chi']]))
 | 
						|
 | 
						|
    plot_values = np.array(plot_values)
 | 
						|
 | 
						|
    column_index = {'smooth': 0, 'damp': 1}
 | 
						|
 | 
						|
    keys = np.unique(plot_values[:, column_index[fix]])
 | 
						|
    names = {'smooth': 'Smoothing', 'damp': 'Damping'}
 | 
						|
 | 
						|
    for key in keys:
 | 
						|
        plot_line = plot_values[plot_values[:, column_index[fix]] == key]
 | 
						|
        second_index = column_index['smooth'] if fix == 'damp' else column_index['damp']
 | 
						|
        plot_line = np.array(sorted(plot_line, key=lambda x: x[second_index]))
 | 
						|
        norms = plot_line[:, 2]
 | 
						|
        chis = plot_line[:, 3]
 | 
						|
        #text = [str(item) for item in plot_line[:, second_index]]
 | 
						|
 | 
						|
        x = np.linspace(min(norms), max(norms), num=100)
 | 
						|
 | 
						|
        #popt, pcov = curve_fit(one_over_x, norms, chis, method='trf')#, bounds=[min(norms), max(norms)])
 | 
						|
        #fit_result = one_over_x(x, *popt)
 | 
						|
        #line = plt.plot(x, fit_result, ':', lw=0.8)[0]
 | 
						|
        fninfo = os.path.split(fname_in)[-1].replace('.json', '').split('_f')[-1]
 | 
						|
        label = '{}: {:g}'.format(names[fix], float(key))
 | 
						|
 | 
						|
        line = plt.plot(norms, chis, linestyle=next(linecycler), lw=0.8, label=label)[0]
 | 
						|
        #coefs = poly.polyfit(norms, chis, 4)
 | 
						|
        #ffit = poly.polyval(x, coefs)
 | 
						|
        #line = plt.plot(x, ffit, ':', lw=0.8)[0]
 | 
						|
 | 
						|
        #label = label='{}: {:g} (smgrad: {})'.format(names[fix], float(key), fninfo)
 | 
						|
        plt.plot(norms, chis, c=line.get_color(), marker='.', lw=0.)
 | 
						|
        #plt.text(norms, chis, text)
 | 
						|
        for item in plot_line:
 | 
						|
            plt.text(item[2], item[3], str(item[second_index]), horizontalalignment='left')
 | 
						|
    #plt.title('Plot of Misfit against Norm ({})'.format(plot_norm))
 | 
						|
 | 
						|
 | 
						|
if __name__ == '__main__':
 | 
						|
    #calc_tradeoff('/data/AlpArray_Data/fmtomo/v5/tradeoff_curves/crust_included_grad_smooth_FIXED_dts_grad_1.5_sm*_damp*/',
 | 
						|
    #             '/data/AlpArray_Data/various/alparray/tradeoff_v5_f1.5.json')
 | 
						|
    #calc_tradeoff('/data/AlpArray_Data/fmtomo/v5/tradeoff_curves/crust_included_grad_smooth_FIXED_dts_sm*_damp*/',
 | 
						|
    #              '/data/AlpArray_Data/various/alparray/tradeoff_v5_f2.0.json')
 | 
						|
 | 
						|
    fig = normed_figure(width_cm=10, ratio=1.)
 | 
						|
    #tradeoff_infiles = ['tradeoff_v4_f1.5.json', 'tradeoff_v4_f3.json', 'tradeoff_v4_f10.json']
 | 
						|
    tradeoff_infiles = ['tradeoff_v5_f2.0.json']#, 'tradeoff_v5_f1.5.json']
 | 
						|
    for infile in tradeoff_infiles:
 | 
						|
        infile = os.path.join('/data/AlpArray_Data/various/alparray/', infile)
 | 
						|
        plot_tradeoff(infile, fix='damp', plot_norm='norm')
 | 
						|
 | 
						|
    plt.xlim([1900, 16200])
 | 
						|
    plt.ylim([2.72, 3.8])
 | 
						|
    plt.xlabel('Norm')
 | 
						|
    #plt.ylabel(r'Misfit($\frac{\chi^2}{N}$)')
 | 
						|
    plt.ylabel(r'Misfit($\chi^2/N$)')
 | 
						|
    #plt.title('Tradeoff curve Misfit vs Norm. Numbers in plot show smoothing values.')
 | 
						|
    plt.legend()
 | 
						|
    #plt.show()
 | 
						|
    plt.savefig('/data/AlpArray_Data/sciebo/AlpArray_home/pictures/paper_II/tradeoff.pdf', dpi=300)
 | 
						|
 |