paleochrono.py

"""
TODO: what about symbolic links in github?
TODO: extend the chronology down to the bedrock by extrapolating the accumulation
TODO: optinally use a restart file to have a bootstrap method
TODO: is there an elegant way to unpack the variables vector in the model function?
TODO: allow to save the correction vector to be able to restart while changing the resolution
TODO: include some checks for when ddelta_depth/dz>1
TODO: Delta-depth observations should be lognormal?
TODO: we should superpose two charts for ice and air ages, one for the age and
    one for the uncertainty, since the min age is not always near 0.
TODO: also compute the prior uncertainties and show them in the figures.
TODO: is there really a computation gain with the change of variable for the
    correction functions? Avoiding this change of variables would make the code
    easier to understand. I think there is no gain since solving A^-1 b when we
    have the LU factorisation of A does not cost more than computing A^-1 * b
    when we have computed A^-1.
TODO: make log plot for sedimentation rate, thinning and LID.
TODO: Divide cost fct by 2 for consistency with least_squares.
"""


import sys
import time
import multiprocessing
import math as m
import numpy as np
import matplotlib.pyplot as mpl
from scipy.optimize import least_squares
from scipy.linalg import solve_triangular
from numpy.linalg import cholesky
from scipy.sparse.linalg import LinearOperator
from scipy import stats
import pccfg
from pcsite import Site
from pcsitepair import SitePair
from functools import partial
import gc
from numpy import dot
import os
if os.name != 'nt':
    import resource
# else:
#     import psutil

# Registration of start time
START_TIME = time.perf_counter()

# Read parameter file
pccfg.read_parameters()

# Opening of output.txt file
OUTPUT_FILE = open(pccfg.datadir+'output.txt', 'w')

# Global
VARIABLES = np.array([])
D = {}
DC = {}


def pcprint(message):
    print(message)
    OUTPUT_FILE.write(message+'\n')
    return


def residuals(var):
    """Calculate the residuals as a function of the variables vector."""
    index = 0
    for i, dlab in enumerate(pccfg.list_sites):
        D[dlab].variables = var[index:index+np.size(D[dlab].variables)]
        index = index+np.size(D[dlab].variables)
        D[dlab].model(D[dlab].variables)
#    pccfg.nb_runs = pccfg.nb_runs + 1
    gc.collect()
    return resid()


def resid():
    """Calculate the residuals without recalculating the model."""
    resi = np.array([])
    for i, dlab in enumerate(pccfg.list_sites):
        resi = np.concatenate((resi, D[dlab].variables))
        resi = np.concatenate((resi, D[dlab].residuals()))
        for j, dlab2 in enumerate(pccfg.list_sites):
            # Note that if I put a new i loop here, to separate the D and DC
            # terms, the model runs slower
            if j < i:
                resi = np.concatenate((resi, DC[dlab2+'-'+dlab].residuals()))
    return resi

def obs_resid():
    """Calculate the observation residuals without recalculating the model."""
    resi = np.array([])
    for i, dlab in enumerate(pccfg.list_sites):
        resi = np.concatenate((resi, D[dlab].residuals()))
        for j, dlab2 in enumerate(pccfg.list_sites):
            # Note that if I put a new i loop here, to separate the D and DC
            # terms, the model runs slower
            if j < i:
                resi = np.concatenate((resi, DC[dlab2+'-'+dlab].residuals()))
    return resi

def prior_resid():
    """Calculate the residuals without recalculating the model."""
    resi = np.array([])
    for i, dlab in enumerate(pccfg.list_sites):
        resi = np.concatenate((resi, D[dlab].variables))
    return resi

def residuals_plot():
    """Plot the histogram of the residuals."""
    
    if pccfg.show_prior_residuals:
        fig, ax1 = mpl.subplots()
        mpl.title('Global residuals')
        mpl.xlabel('Residuals (no unit)')
        mpl.ylabel('Probability density')
        resi = resid()
        rms = m.sqrt(np.sum(resi**2)/len(resi))
        mini = np.min(resi, initial=0)
        maxi = np.max(resi, initial=0)
        student = stats.t.fit(resi)
        mpl.hist(resi, bins=40, range=(-4., 4.), density=True, 
                 label=f"RMS: {rms:.3}, min: {mini:.3}, max: {maxi:.3},\n"
                 f"loc: {student[1]:.3}, scale: {student[2]:.3}, df: {student[0]:.3e}")
        x_low, x_up, y_low, y_up = mpl.axis()
        mpl.axis((-4., 4., y_low, y_up))
        mpl.legend()
        mpl.savefig(pccfg.datadir+'/residuals.'+pccfg.fig_format,
                    format=pccfg.fig_format, bbox_inches='tight')
        if not pccfg.show_figures:
            mpl.close()

    if pccfg.show_prior_residuals:
        fig, ax1 = mpl.subplots()
        mpl.title('Prior residuals')
        mpl.xlabel('Residuals (no unit)')
        mpl.ylabel('Probability density')
        resi = prior_resid()
        rms = m.sqrt(np.sum(resi**2)/len(resi))
        mini = np.min(resi, initial=0)
        maxi = np.max(resi, initial=0)
        student = stats.t.fit(resi)
        mpl.hist(resi, bins=40, range=(-4., 4.), density=True, 
                 label=f"RMS: {rms:.3}, min: {mini:.3}, max: {maxi:.3},\n"
                 f"loc: {student[1]:.3}, scale: {student[2]:.3}, df: {student[0]:.3e}")
        x_low, x_up, y_low, y_up = mpl.axis()
        mpl.axis((-4., 4., y_low, y_up))
        mpl.legend()
        mpl.savefig(pccfg.datadir+'/prior_residuals.'+pccfg.fig_format,
                    format=pccfg.fig_format, bbox_inches='tight')
        if not pccfg.show_figures:
            mpl.close()

    fig, ax1 = mpl.subplots()
    mpl.title('Observation residuals')
    mpl.xlabel('Residuals (no unit)')
    mpl.ylabel('Probability density')
    resi = obs_resid()
    rms = m.sqrt(np.sum(resi**2)/len(resi))
    mini = np.min(resi, initial=0)
    maxi = np.max(resi, initial=0)
    student = stats.t.fit(resi)
    mpl.hist(resi, bins=40, range=(-4., 4.), density=True, 
             label=f"RMS: {rms:.3}, min: {mini:.3}, max: {maxi:.3},\n"
             f"loc: {student[1]:.3}, scale: {student[2]:.3}, df: {student[0]:.3e}")
    x_low, x_up, y_low, y_up = mpl.axis()
    mpl.axis((-4., 4., y_low, y_up))
    mpl.legend()
    mpl.savefig(pccfg.datadir+'/obs_residuals.'+pccfg.fig_format,
                format=pccfg.fig_format, bbox_inches='tight')
    if not pccfg.show_figures:
        mpl.close()
        
def cost_function(var):
    """Calculate the cost function terms related to a pair of sites."""
    res = residuals(var)
    cost = dot(res, np.transpose(res)) / 2.
    return cost


def jacob_column(resizero, dlabj, k):
    delta = m.sqrt(np.finfo(float).eps)  # Stolen from the leastsq code
    D[dlabj].variables[k] += delta
    D[dlabj].model(D[dlabj].variables)
    deriv = [np.array([])]
    index = 0
    for i, dlab in enumerate(pccfg.list_sites):
        index = index+len(D[dlab].variables)
        if dlabj == dlab:
            der = np.zeros(len(D[dlab].variables))
            der[k] = 1.
            deriv.append(der)
            der = (D[dlab].residuals() -
                   resizero[index:index+RESI_SIZE[i, i]]) / delta
            deriv.append(der)
        else:
            deriv.append(np.zeros(len(D[dlab].variables)))
            deriv.append(np.zeros(RESI_SIZE[i, i]))
        index = index+RESI_SIZE[i, i]
        for j, dlab2 in enumerate(pccfg.list_sites):
            if j < i:
                if dlabj == dlab or dlabj == dlab2:
                    der = (DC[dlab2 + '-' + dlab].residuals() -
                           resizero[index:index+RESI_SIZE[j, i]])/delta
                    deriv.append(der)
                else:
                    deriv.append(np.zeros(RESI_SIZE[j, i]))
                index = index+RESI_SIZE[j, i]
    D[dlabj].variables[k] -= delta
    return np.concatenate(deriv)


def jacobian_analytical(var):
    """Calculate the Jacobian of each residual term
    with analytical formulas."""
    jac_list = []
    for k, dlabj in enumerate(pccfg.list_sites):
        D[dlabj].corrected_jacobian()
        deriv = []
        for i, dlab in enumerate(pccfg.list_sites):
            if dlabj == dlab:
                deriv.append(np.diag(np.ones(len(D[dlab].variables))))
                deriv.append(D[dlab].residuals_jacobian())
            else:
                deriv.append(np.zeros((len(D[dlabj].variables),
                                       len(D[dlab].variables))))
                deriv.append(np.zeros((len(D[dlabj].variables),
                                       RESI_SIZE[i, i])))
            for j, dlab2 in enumerate(pccfg.list_sites):
                if j < i:
                    if dlabj == dlab:
                        deriv.append(DC[dlab2+'-'+dlab].residuals_jacobian2())
                    elif dlabj == dlab2:
                        deriv.append(DC[dlab2+'-'+dlab].residuals_jacobian1())
                    else:
                        deriv.append(np.zeros((len(D[dlabj].variables),
                                               RESI_SIZE[j, i])))
        jac_list.append(np.concatenate(deriv, axis=1))
    jacob = np.concatenate(jac_list)
#    print(np.shape(jacob), np.shape(resid()), len(VARIABLES))
    return np.transpose(jacob)


def jacobian_semi_adjoint(var):

    jac = np.array([[None for _ in range(len(pccfg.list_sites))]
                    for _ in range(len(pccfg.list_sites))])
    for i, dlab in enumerate(pccfg.list_sites):
        D[dlab].corrected_jacobian()
        for j, dlab2 in enumerate(pccfg.list_sites):
            if j == i:
                jac[i, i] = D[dlab].residuals_jacobian()
            if j < i:
                jac[j, i] = DC[dlab2+'-'+dlab].residuals_jacobian2()
                jac[i, j] = DC[dlab2+'-'+dlab].residuals_jacobian1()

    def mv(v):

        index = 0
        resi = np.array([])
        for i, dlab in enumerate(pccfg.list_sites):
            # Why do we need to sometimes flatten here? Strange.
            D[dlab].var_delta = v[index:index+np.size(D[dlab].variables)]\
                .flatten()
            index = index+np.size(D[dlab].variables)
        for i, dlab in enumerate(pccfg.list_sites):
            # Why do we need to sometimes flatten here? Strange.
            resi = np.concatenate((resi, D[dlab].var_delta))
            resi = np.concatenate((resi, dot(np.transpose(jac[i, i]),
                                                D[dlab].var_delta)))
            for j, dlab2 in enumerate(pccfg.list_sites):
                # Note that if I put a new i loop here,
                # to separate the D and DC terms, the model runs slower
                if j < i:
                    resi = np.concatenate((resi,
                                           dot(np.transpose(jac[j, i]),
                                                  D[dlab].var_delta) +
                                           dot(np.transpose(jac[i, j]),
                                                  D[dlab2].var_delta)))

        return resi

    def rmv(v):

        vari = []
        for k, dlabj in enumerate(pccfg.list_sites):
            vari = vari + [np.zeros(np.size(D[dlabj].variables))]

        index = 0
        for i, dlab in enumerate(pccfg.list_sites):
            vari[i] = v[index:index+np.size(D[dlab].variables)].flatten()
            index = index+np.size(D[dlab].variables)
            vari[i] = vari[i] +\
                dot(jac[i, i], v[index:index+RESI_SIZE[i, i]])
            # vari[i] = vari[i] +
            # D[dlab].residuals_adj( v[index:index+RESI_SIZE[i,i]])
            index = index+RESI_SIZE[i, i]
            for j, dlab2 in enumerate(pccfg.list_sites):
                if j < i:
                    vari[i] = vari[i] + dot(jac[j, i],
                                               v[index:index+RESI_SIZE[j, i]])
                    vari[j] = vari[j]+dot(jac[i, j],
                                             v[index:index+RESI_SIZE[j, i]])
                    index = index + RESI_SIZE[j, i]

        vari = np.concatenate(vari)

        return vari
#        return dot(np.transpose(jac), v)

    return LinearOperator((RESI_SIZE_TOT, VAR_SIZE), matvec=mv, rmatvec=rmv)


def jacobian_adjoint(var):
    """Full adjoint method. Not ready yet."""
    # FIXME: Adjoint give slightly more iterations than semi_adjoint
    # on the med exp.
    # Check what is the issue.
    print('Full adjoint is not ready yet. Exiting.')
    sys.exit()

    jac = np.array([[None for _ in range(len(pccfg.list_sites))]
                    for _ in range(len(pccfg.list_sites))])
    for i, dlab in enumerate(pccfg.list_sites):
        D[dlab].corrected_jacobian()
        for j, dlab2 in enumerate(pccfg.list_sites):
            if j == i:
                jac[i, i] = D[dlab].residuals_jacobian()
            if j < i:
                jac[j, i] = DC[dlab2+'-'+dlab].residuals_jacobian2()
                jac[i, j] = DC[dlab2+'-'+dlab].residuals_jacobian1()

    def mv(v):

        index = 0
        resi = np.array([])
        for i, dlab in enumerate(pccfg.list_sites):
            # Why do we need to sometimes flatten here? Strange.
            D[dlab].var_delta = v[index:index+np.size(D[dlab].variables)]\
                .flatten()
            index = index+np.size(D[dlab].variables)
            resi = np.concatenate((resi, D[dlab].var_delta))
            D[dlab].model_delta(D[dlab].var_delta)
            resi = np.concatenate((resi, D[dlab].residuals_delta()))
            for j, dlab2 in enumerate(pccfg.list_sites):
                # Note that if I put a new i loop here,
                # to separate the D and DC terms, the model runs slower
                if j < i:
                    resi = np.concatenate((resi,
                                           DC[dlab2 + '-' +
                                              dlab].residuals_delta()))
        return resi

    def rmv(v):

        vari =[]
        for k, dlabj in enumerate(pccfg.list_sites):
            vari = vari + [np.zeros(np.size(D[dlabj].variables))]

        index = 0
        for i, dlab in enumerate(pccfg.list_sites):
            vari[i] = v[index:index+np.size(D[dlab].variables)].flatten()
            index = index+np.size(D[dlab].variables)
            vari[i] = vari[i] + dot(jac[i,i], v[index:index+RESI_SIZE[i,i]])
#            vari[i] = vari[i] + D[dlab].residuals_adj( v[index:index+RESI_SIZE[i,i]])
            index = index+RESI_SIZE[i,i]
            for j, dlab2 in enumerate(pccfg.list_sites):
                if j < i:
                    vari[i] = vari[i]+dot(jac[j,i],
                        v[index:index+RESI_SIZE[j,i]])
                    vari[j] = vari[j]+dot(jac[i,j],
                        v[index:index+RESI_SIZE[j,i]])
                    index = index + RESI_SIZE[j,i]
        
        vari = np.concatenate(vari)

        return vari        
#        return dot(np.transpose(jac), v)
    
    return LinearOperator((RESI_SIZE_TOT, VAR_SIZE), matvec=mv, rmatvec=rmv)

def jacobian_semi_analytical(var):
    """Calculate the Jacobian with a finite difference scheme for each block."""
    resizero = residuals(var)
    jac_list = []
    for k, dlabj in enumerate(pccfg.list_sites):
        if pccfg.is_parallel:
            list_args = list(range(len(D[dlabj].variables)))
            if __name__ == "__main__":
                with multiprocessing.Pool(pccfg.nb_nodes) as pool:
                    results = pool.map(partial(jacob_column, resizero, dlabj),
                                               list_args)
                jac_list.append(results)
        else:
            for l in range(len(D[dlabj].variables)):
#                jacob = np.vstack((jacob, jacob_column(resizero, dlabj, l)))
                jac_list.append(np.array([jacob_column(resizero, dlabj, l)]))
        D[dlabj].model(D[dlabj].variables)
    jacob = np.concatenate(jac_list)
    return np.transpose(jacob)

def jacobian_numerical(var):
    """Calculate the Jacobian with a finite difference scheme."""
    zeropred = residuals(var)
    derivparams = []
    results = []
    delta = m.sqrt(np.finfo(float).eps) #Stolen from the leastsq code
    #fixme: This loop is probably sub-optimal. Have a look at what does leastsq to improve this.
#        results.append(residuals(derivparams))
    if pccfg.is_parallel:
        for i in range(len(var)):
            copy = np.array(var)
            copy[i] += delta
            derivparams.append(copy)
        if __name__ == "__main__":
            pool = multiprocessing.Pool(pccfg.nb_nodes)
        results = pool.map(residuals, derivparams)
        derivs = [(r - zeropred)/delta for r in results]
    else:
        list_derivs = []
        for i in range(len(var)):
            copy = np.array(var)
            copy[i] += delta
            list_derivs.append(np.array([(residuals(copy)-zeropred)/delta]))
        derivs = np.concatenate(list_derivs)
    return np.transpose(derivs)

##MAIN

##Initialisation
RESI_SIZE = np.empty((np.size(pccfg.list_sites), np.size(pccfg.list_sites)), dtype=int)

for di, dlabel in enumerate(pccfg.list_sites):

    print('Initialization of site '+dlabel)

    D[dlabel] = Site(dlabel)
    D[dlabel].model(D[dlabel].variables)
#    D[dlabel].a_init=D[dlabel].a
#    D[dlabel].lid_init=D[dlabel].lid
    D[dlabel].write_init()
#    D[dlabel].display_init()
    VARIABLES = np.concatenate((VARIABLES, D[dlabel].variables))
    RESI_SIZE[di, di] = np.size(D[dlabel].residuals())

for di, dlabel in enumerate(pccfg.list_sites):
    for dj, dlabel2 in enumerate(pccfg.list_sites):
        if dj < di:
            print('Initialization of site pair '+dlabel2+'-'+dlabel)
            DC[dlabel2+'-'+dlabel] = SitePair(D[dlabel2], D[dlabel])
#            DC[dlabel2+'-'+dlabel].display_init()
            RESI_SIZE[dj, di] = np.size(DC[dlabel2+'-'+dlabel].residuals())

VAR_SIZE = len(VARIABLES)
RESI_SIZE_TOT = len(resid())
pcprint('Size of VARIABLES vector: '+ str(VAR_SIZE))
pcprint('Size of RESIDUALS vector: '+ str(RESI_SIZE_TOT))

##Optimization
START_TIME_OPT = time.perf_counter()
pcprint('Initial cost function: '+ str(cost_function(VARIABLES)))
#print(jacobian_semi_analytical(VARIABLES))
#print(jacobian_analytical(VARIABLES))
if pccfg.opt_method == 'leastsq':
    pccfg.opt_method = 'trf'
    pccfg.is_parallel = False
elif pccfg.opt_method == 'leastsq-parallel':
    pccfg.opt_method = 'trf'
    pccfg.is_parallel = True
if pccfg.opt_method == "trf" or pccfg.opt_method == 'lm':
    print('Optimization by:', pccfg.opt_method)
    if pccfg.opt_method == 'trf':
        print('tr_solver:', pccfg.tr_solver)
    print('Jabobian:', pccfg.jacobian)
    if pccfg.jacobian == 'automatic':
        jac = '2-point'
    else:
        jac = eval('jacobian_'+pccfg.jacobian)
    OptimizeResult = least_squares(residuals, VARIABLES, method=pccfg.opt_method,
                                   jac=jac,
                                   tr_solver=pccfg.tr_solver,
                                   xtol=pccfg.tol, ftol=pccfg.tol, gtol=pccfg.tol, verbose=2)
    VARIABLES = OptimizeResult.x
    pcprint('Optimized cost function: ' + str(cost_function(VARIABLES)))
    pcprint('Optimization execution time: ' + str(time.perf_counter() - START_TIME_OPT) + ' seconds')
    if pccfg.jacobian == 'adjoint' or pccfg.jacobian == 'semi_adjoint':
        print('Calculating Jacobian matrix.')
        JACMAT = jacobian_analytical(VARIABLES)
#        print('Size of JACMAT in kbytes',JACMAT.nbytes/1000)
        for dlabel in pccfg.list_sites:
            D[dlabel].corrected_jacobian_free()
    else:
        JACMAT = OptimizeResult.jac
    print('Calculating Hessian matrix.')
    HESS = dot(np.transpose(JACMAT), JACMAT)
    JACMAT = None
elif pccfg.opt_method == 'none':
    print('No optimization')
    VARIABLES = np.zeros(np.size(VARIABLES))
    HESS = np.diag(np.ones(np.size(VARIABLES)))
else:
    print(pccfg.opt_method, ': Optimization method not recognized.')
    sys.exit()
#print 'solution: ',VARIABLES

print('Factorisation of the Hessian matrix')
HESS_chol = cholesky(HESS)
HESS = None
#HESS_chol = cholesky(HESS, overwrite_a=True, check_finite=False)
#HESS_chol = np.transpose(HESS_chol)
#COV = np.linalg.inv(HESS)
#HESS = None

print('Calculation of confidence intervals')
# COV = np.linalg.inv(HESS)
INDEXSITE = 0
for dlabel in pccfg.list_sites:
    print('Covariance matrix for '+dlabel)
#    input('Before solving the triangular system. Program paused.')
    SIZESITE = np.size(D[dlabel].variables)
    D[dlabel].variables = VARIABLES[INDEXSITE:INDEXSITE+SIZESITE]
    block1 = np.zeros((INDEXSITE, SIZESITE))
    block2 = np.diag(np.ones(SIZESITE))
    block3 = np.zeros((np.size(VARIABLES)-INDEXSITE-SIZESITE, SIZESITE))
    block = np.vstack((block1, block2, block3))
    toto = solve_triangular(HESS_chol, block, lower=True)
    block = None
    D[dlabel].cov = dot(np.transpose(toto), toto)
    toto = None
#    D[dlabel].cov = COV[INDEXSITE:INDEXSITE+SIZESITE,INDEXSITE:INDEXSITE+SIZESITE]
    INDEXSITE = INDEXSITE+np.size(D[dlabel].variables)
#    input('Before calculating sigma. Program paused.')
#    COV[INDEXSITE:INDEXSITE+SIZESITE,:] = None
#    COV[:,INDEXSITE:INDEXSITE+SIZESITE] = None
HESS_chol = None
for dlabel in pccfg.list_sites:
    print('Confidence intervals for '+dlabel)
    if pccfg.calc_errors:
        D[dlabel].sigma()
    D[dlabel].cov = None

# Final display and output
print('Figures, outliers detection and saving')
for di, dlabel in enumerate(pccfg.list_sites):
    print('Figures, outliers detection and saving of', dlabel)
    D[dlabel].save()
    D[dlabel].figures()
for di, dlabel in enumerate(pccfg.list_sites):
    for dj, dlabel2 in enumerate(pccfg.list_sites):
        if dj < di:
            print('Display of '+dlabel2+'-'+dlabel+' site pair')
            DC[dlabel2+'-'+dlabel].figures()

# Plotting histogram of residuals
residuals_plot()

print('Checking for air age inversion')
for di, dlabel in enumerate(pccfg.list_sites):
    if D[dlabel].check_airage_inversion():
        print("WARNING: air age inversion for", dlabel)

# Program execution time
pcprint('Program execution time: '+str(time.perf_counter()-START_TIME)+' seconds')
if os.name != 'nt':
    mem = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    pcprint('Max memory usage: '+str(mem)+' kbytes')
# else:
#     process = psutil.Process(os.getpid())
#     mem = int(process.memory_info().rss / 1024)

if pccfg.show_figures:
    mpl.show()

# Closing output file
OUTPUT_FILE.close()