Mentions légales du service

Skip to content
Snippets Groups Projects
fedic.py 28.9 KiB
Newer Older
Martin Genet's avatar
Martin Genet committed
#coding=utf8

########################################################################
###                                                                  ###
### Created by Martin Genet, 2016                                    ###
###                                                                  ###
### École Polytechnique, Palaiseau, France                           ###
###                                                                  ###
########################################################################

import dolfin
import glob
Martin Genet's avatar
Martin Genet committed
import numpy
import os
Martin Genet's avatar
Martin Genet committed

import myVTKPythonLibrary as myVTK

Martin Genet's avatar
Martin Genet committed
import myFEniCSPythonLibrary as myFEniCS
Martin Genet's avatar
Martin Genet committed
from print_tools import *
Martin Genet's avatar
Martin Genet committed

########################################################################

Martin Genet's avatar
Martin Genet committed
dolfin.parameters["form_compiler"]["optimize"] = False # can't use that for "complex" mechanical models…
dolfin.parameters["form_compiler"]["cpp_optimize"] = True
# dolfin.parameters["num_threads"] = 8 # 2016-07-07: doesn't seem to work…

########################################################################
Martin Genet's avatar
Martin Genet committed

Martin Genet's avatar
Martin Genet committed
def fedic(
        working_folder,
        working_basename,
        images_folder,
        images_basename,
Martin Genet's avatar
Martin Genet committed
        images_zfill=2,
        images_n_frames=None,
        images_ref_frame=0,
        images_expressions_type="cpp", # cpp, py
        images_dynamic_scaling=0,
        mesh=None,
        mesh_folder=None,
        mesh_basename=None,
Martin Genet's avatar
Martin Genet committed
        mesh_degree=1,
        regul_type="neo-hookean",
        regul_level=0.1,
        regul_poisson=0.0,
        tangent_type="Idef", # Idef, Idef-wHess, Iold, Iref
        residual_type="Iref", # Iref, Iold, Iref-then-Iold
Martin Genet's avatar
Martin Genet committed
        relax_type="constant", # constant, aitken, gss
Martin Genet's avatar
Martin Genet committed
        relax_init=1.0,
Martin Genet's avatar
Martin Genet committed
        tol_res=None,
Martin Genet's avatar
Martin Genet committed
        n_iter_max=100,
        continue_after_fail=0,
        print_iterations=0):
Martin Genet's avatar
Martin Genet committed

Martin Genet's avatar
Martin Genet committed
    tab = 0
Martin Genet's avatar
Martin Genet committed

    print_str(tab,"Checking number of frames…")
    if (images_n_frames is None):
        images_n_frames = len(glob.glob(images_folder+"/"+images_basename+"_"+"[0-9]"*images_zfill+".vti"))
    assert (images_n_frames > 1), "images_n_frames = "+str(images_n_frames)+" <= 1. Aborting."
    print_var(tab+1,"images_n_frames",images_n_frames)

    assert (abs(images_ref_frame) < images_n_frames), "abs(images_ref_frame) = "+str(images_ref_frame)+" >= images_n_frames. Aborting."
    images_ref_frame = images_ref_frame%images_n_frames
    print_var(tab+1,"images_ref_frame",images_ref_frame)
    print_str(tab,"Loading mesh…")
    assert (mesh is not None or ((mesh_folder is not None) and (mesh_basename is not None))), "Must provide a mesh (mesh = "+str(mesh)+") or a mesh file (mesh_folder = "+str(mesh_folder)+", mesh_basename = "+str(mesh_basename)+"). Aborting."
    if (mesh is None):
        mesh_filebasename = mesh_folder+"/"+mesh_basename
        mesh_filename = mesh_filebasename+"."+"xml"
        assert os.path.exists(mesh_filename), "No mesh in "+mesh_filename+". Aborting."
        mesh = dolfin.Mesh(mesh_filename)
    dX = dolfin.dx(mesh)
Martin Genet's avatar
Martin Genet committed
    mesh_V0 = dolfin.assemble(dolfin.Constant(1)*dX)
    print_var(tab+1,"mesh_n_cells",len(mesh.cells()))
    print_sci(tab+1,"mesh_V0",mesh_V0)
Martin Genet's avatar
Martin Genet committed
    print_str(tab,"Computing quadrature degree for images…")
    ref_image_filename = images_folder+"/"+images_basename+"_"+str(images_ref_frame).zfill(images_zfill)+".vti"
        images_quadrature = myFEniCS.compute_quadrature_degree_from_points_count(
            mesh_filebasename=mesh_filebasename,
            verbose=1)
Martin Genet's avatar
Martin Genet committed
    print_var(tab+1,"images_quadrature",images_quadrature)
Martin Genet's avatar
Martin Genet committed

Martin Genet's avatar
Martin Genet committed
    print_str(tab,"Loading reference image…")
Martin Genet's avatar
Martin Genet committed
    images_dimension = myVTK.computeImageDimensionality(
Martin Genet's avatar
Martin Genet committed
        verbose=0)
    print_var(tab+1,"images_dimension",images_dimension)
Martin Genet's avatar
Martin Genet committed
    assert (images_dimension in (2,3)), "images_dimension must be 2 or 3. Aborting."
Martin Genet's avatar
Martin Genet committed
    fe = dolfin.FiniteElement(
        family="Quadrature",
Martin Genet's avatar
Martin Genet committed
        degree=images_quadrature,
        quad_scheme="default")
    ve = dolfin.VectorElement(
        family="Quadrature",
        cell=mesh.ufl_cell(),
Martin Genet's avatar
Martin Genet committed
        degree=images_quadrature,
    te = dolfin.TensorElement(
        family="Quadrature",
        cell=mesh.ufl_cell(),
        degree=images_quadrature,
        quad_scheme="default")
    te._quad_scheme = "default"                       # shouldn't be needed
    for k in xrange(images_dimension**2):             # shouldn't be needed
        te.sub_elements()[k]._quad_scheme = "default" # shouldn't be needed
    if (images_expressions_type == "cpp"):
        Iref = dolfin.Expression(
            cppcode=myFEniCS.get_ExprIm_cpp(
                im_dim=images_dimension,
                im_type="im",
                im_is_def=0),
Martin Genet's avatar
Martin Genet committed
            element=fe)
        Iref.init_image(ref_image_filename)
        DIref = dolfin.Expression(
            cppcode=myFEniCS.get_ExprIm_cpp(
                im_dim=images_dimension,
                im_type="grad",
                im_is_def=0),
        DIref.init_image(ref_image_filename)
    elif (images_expressions_type == "py"):
        if (images_dimension == 2):
            Iref = myFEniCS.ExprIm2(
                filename=ref_image_filename,
                element=fe)
            DIref = myFEniCS.ExprGradIm2(
                filename=ref_image_filename,
                element=ve)
        elif (images_dimension == 3):
            Iref = myFEniCS.ExprIm3(
                filename=ref_image_filename,
                element=fe)
            DIref = myFEniCS.ExprGradIm3(
                filename=ref_image_filename,
                element=ve)
    else:
        assert (0), "\"images_expressions_type\" (="+str(images_expressions_type)+") must be \"cpp\" or \"py\". Aborting."
    Iref_int = dolfin.assemble(Iref * dX, form_compiler_parameters={'quadrature_degree':images_quadrature})/mesh_V0
    Iref_norm = (dolfin.assemble(Iref**2 * dX, form_compiler_parameters={'quadrature_degree':images_quadrature})/mesh_V0)**(1./2)
Martin Genet's avatar
Martin Genet committed
    assert (Iref_norm > 0.), "Iref_norm = "+str(Iref_norm)+" <= 0. Aborting."
    print_var(tab+1,"Iref_int",Iref_int)
Martin Genet's avatar
Martin Genet committed
    print_var(tab+1,"Iref_norm",Iref_norm)
Martin Genet's avatar
Martin Genet committed

Martin Genet's avatar
Martin Genet committed
    print_str(tab,"Defining functions…")
Martin Genet's avatar
Martin Genet committed
    vfs = dolfin.VectorFunctionSpace(
        mesh=mesh,
        family="Lagrange",
        degree=mesh_degree)
Martin Genet's avatar
Martin Genet committed
    U = dolfin.Function(
Martin Genet's avatar
Martin Genet committed
        vfs,
Martin Genet's avatar
Martin Genet committed
        name="displacement")
    U.vector().zero()
Martin Genet's avatar
Martin Genet committed
        vfs,
        name="previous displacement")
    Uold.vector().zero()
Martin Genet's avatar
Martin Genet committed
    dU = dolfin.Function(
Martin Genet's avatar
Martin Genet committed
        vfs,
        name="displacement correction")
Martin Genet's avatar
Martin Genet committed
    dU_ = dolfin.TrialFunction(vfs)
    dV_ = dolfin.TestFunction(vfs)
Martin Genet's avatar
Martin Genet committed

    print_str(tab,"Printing initial solution…")
    if not os.path.exists(working_folder):
        os.mkdir(working_folder)
    pvd_basename = working_folder+"/"+working_basename+"_"
    file_pvd = dolfin.File(pvd_basename+".pvd")
    for vtu_filename in glob.glob(pvd_basename+"*.vtu"):
        os.remove(vtu_filename)
    file_pvd << (U, float(images_ref_frame))

    if (print_iterations):
        for filename in glob.glob(working_folder+"/"+working_basename+"-frame=*.*"):
            os.remove(filename)

    print_str(tab,"Defining mechanical energy…")
    E     = dolfin.Constant(1.0)
    nu    = dolfin.Constant(regul_poisson)
    kappa = E/3/(1-2*nu)         # = E/3 if nu = 0
    lmbda = E*nu/(1+nu)/(1-2*nu) # = 0   if nu = 0
    mu    = E/2/(1+nu)           # = E/2 if nu = 0
    C1    = mu/2
    C2    = mu/2
    D1    = kappa/2

Martin Genet's avatar
Martin Genet committed
    I = dolfin.Identity(images_dimension)
    F = I + dolfin.grad(U)
    J = dolfin.det(F)
    if   (regul_type == "laplacian"): # <- super bad
        e     = dolfin.sym(dolfin.grad(U))
        psi_m = (lmbda * dolfin.tr(e)**2 + 2*mu * dolfin.tr(e*e))/2
    elif (regul_type == "kirchhoff"): # <- pretty bad too
        C     = F.T * F
        E     = (C - I)/2
        psi_m = (lmbda * dolfin.tr(E)**2 + 2*mu * dolfin.tr(E*E))/2
    elif (regul_type == "neo-hookean"):
        C     = F.T * F
        Ic    = dolfin.tr(C)
        Ic0   = dolfin.tr(I)
        psi_m = C1 * (Ic - Ic0 - 2*dolfin.ln(J)) + D1 * (J**2 - 1 - 2*dolfin.ln(J))
    elif (regul_type == "mooney-rivlin"):
        C     = F.T * F
        Ic    = dolfin.tr(C)
        Ic0   = dolfin.tr(I)
        IIc   = (dolfin.tr(C)**2 - dolfin.tr(C*C))/2
        IIc0  = (dolfin.tr(I)**2 - dolfin.tr(I*I))/2
        psi_m = (C1/2) * (Ic - Ic0 - 2*dolfin.ln(J)) + (C2/2) * (IIc - IIc0 - 4*dolfin.ln(J)) + D1 * (J**2 - 1 - 2*dolfin.ln(J))
    else:
        assert (0), "\"regul_type\" must be \"laplacian\", \"kirchhoff\", \"neo-hookean\", or \"mooney-rivlin\". Aborting."

    Dpsi_m  = dolfin.derivative( psi_m, U, dV_)
    DDpsi_m = dolfin.derivative(Dpsi_m, U, dU_)

Martin Genet's avatar
Martin Genet committed
    mesh_V = dolfin.assemble(J*dX)
    print_sci(tab+1,"mesh_V",mesh_V)

    volume_basename = working_folder+"/"+working_basename+"-volume"
    file_volume = open(volume_basename+".dat","w")
    file_volume.write("#k_frame mesh_V"+"\n")
    file_volume.write(" ".join([str(val) for val in [images_ref_frame, mesh_V]])+"\n")

    print_str(tab,"Defining deformed image…")
    scaling = numpy.array([1.,0.])
    #scaling = dolfin.Constant([1.,0.])
    if (images_expressions_type == "cpp"):
        Idef = dolfin.Expression(
            cppcode=myFEniCS.get_ExprIm_cpp(
                im_dim=images_dimension,
                im_type="im",
                im_is_def=1),
Martin Genet's avatar
Martin Genet committed
            element=fe)
        Idef.init_disp(U)
        Idef.init_dynamic_scaling(scaling)
        DIdef = dolfin.Expression(
            cppcode=myFEniCS.get_ExprIm_cpp(
                im_dim=images_dimension,
                im_type="grad",
                im_is_def=1),
        DIdef.init_disp(U)
        DIdef.init_dynamic_scaling(scaling)
            assert (0), "ToDo"
        Iold = dolfin.Expression(
            cppcode=myFEniCS.get_ExprIm_cpp(
                im_dim=images_dimension,
                im_type="im",
                im_is_def=1),
        Iold.init_disp(Uold)
        Iold.init_dynamic_scaling(scaling)
        DIold = dolfin.Expression(
            cppcode=myFEniCS.get_ExprIm_cpp(
                im_dim=images_dimension,
                im_type="grad",
                im_is_def=1),
        DIold.init_disp(Uold)
        DIold.init_dynamic_scaling(scaling)
    elif (images_expressions_type == "py"):
        if (images_dimension == 2):
            Idef = myFEniCS.ExprDefIm2(
                scaling=scaling,
                element=fe)
            DIdef = myFEniCS.ExprGradDefIm2(
                U=U,
                scaling=scaling,
                element=ve)
            if ("-wHess" in tangent_type):
                DDIdef = myFEniCS.ExprHessDefIm2(
                    U=U,
                    scaling=scaling,
                    element=te)
            Iold = myFEniCS.ExprDefIm2(
                U=Uold,
                scaling=scaling,
                element=fe)
            DIold = myFEniCS.ExprGradDefIm2(
                U=Uold,
                scaling=scaling,
                element=ve)
        elif (images_dimension == 3):
            Idef = myFEniCS.ExprDefIm3(
                U=U,
                scaling=scaling,
                element=fe)
            DIdef = myFEniCS.ExprGradDefIm3(
                U=U,
                scaling=scaling,
                element=ve)
            if ("-wHess" in tangent_type):
                DDIdef = myFEniCS.ExprHessDefIm3(
                    U=U,
                    scaling=scaling,
                    element=te)
            Iold = myFEniCS.ExprDefIm3(
                U=Uold,
                scaling=scaling,
                element=fe)
            DIold = myFEniCS.ExprGradDefIm3(
                U=Uold,
                scaling=scaling,
                element=ve)
    else:
        assert (0), "\"images_expressions_type\" (="+str(images_expressions_type)+") must be \"cpp\" or \"py\". Aborting."

    print_str(tab,"Defining correlation energy…")
    psi_c   = (Idef - Iref)**2/2
    Dpsi_c  = (Idef - Iref) * dolfin.dot(DIdef, dV_)
    DDpsi_c = dolfin.dot(DIdef, dU_) * dolfin.dot(DIdef, dV_)
    if ("-wHess" in tangent_type):
        DDpsi_c += (Idef - Iref) * dolfin.inner(dolfin.dot(DDIdef, dU_), dV_)
    Dpsi_c_old  = (Idef - Iold) * dolfin.dot(DIdef, dV_)
    DDpsi_c_old = dolfin.dot(DIold, dU_) * dolfin.dot(DIold, dV_)
    DDpsi_c_ref = dolfin.dot(DIref, dU_) * dolfin.dot(DIref, dV_)

    b0 = Iref * dolfin.dot(DIref, dV_) * dX
    B0 = dolfin.assemble(b0, form_compiler_parameters={'quadrature_degree':images_quadrature})
    res_norm0 = B0.norm("l2")
    assert (res_norm0 > 0.), "res_norm0 = "+str(res_norm0)+" <= 0. Aborting."
    print_var(tab+1,"res_norm0",res_norm0)
Martin Genet's avatar
Martin Genet committed

    #regul_level = dolfin.Constant(regul_level) * 5*(1+nu)*(1-2*nu)/(5-4*nu)
    regul_level = dolfin.Constant(regul_level)

Martin Genet's avatar
Martin Genet committed
    A = None
Martin Genet's avatar
Martin Genet committed
    if (tangent_type == "Iref"):
        print_str(tab,"Assembly…")
        A = dolfin.assemble((1.-regul_level) * DDpsi_c_ref * dX, tensor=A, form_compiler_parameters={'quadrature_degree':images_quadrature})
        A = dolfin.assemble((   regul_level) * DDpsi_m     * dX, tensor=A, add_values=True)
Martin Genet's avatar
Martin Genet committed
    B = None
Martin Genet's avatar
Martin Genet committed

Martin Genet's avatar
Martin Genet committed
    print_str(tab,"Looping over frames…")
Martin Genet's avatar
Martin Genet committed
    n_iter_tot = 0
Martin Genet's avatar
Martin Genet committed
    global_success = True
    for forward_or_backward in ["forward", "backward"]:
Martin Genet's avatar
Martin Genet committed
        print_var(tab,"forward_or_backward",forward_or_backward)
Martin Genet's avatar
Martin Genet committed

        if (forward_or_backward == "forward"):
            k_frames_old = range(images_ref_frame  , images_n_frames-1, +1)
            k_frames     = range(images_ref_frame+1, images_n_frames  , +1)
Martin Genet's avatar
Martin Genet committed
        elif (forward_or_backward == "backward"):
            k_frames_old = range(images_ref_frame  ,  0, -1)
            k_frames     = range(images_ref_frame-1, -1, -1)
Martin Genet's avatar
Martin Genet committed
        print_var(tab,"k_frames",k_frames)
Martin Genet's avatar
Martin Genet committed

        if (forward_or_backward == "backward"):
            U.vector().zero()
            Uold.vector().zero()
Martin Genet's avatar
Martin Genet committed

Martin Genet's avatar
Martin Genet committed
        tab += 1
        success = True
        for (k_frame,k_frame_old) in zip(k_frames,k_frames_old):
Martin Genet's avatar
Martin Genet committed
            print_var(tab-1,"k_frame",k_frame)
Martin Genet's avatar
Martin Genet committed

            if (print_iterations):
Martin Genet's avatar
Martin Genet committed
                frame_basename = working_folder+"/"+working_basename+"-frame="+str(k_frame).zfill(images_zfill)
                file_dat_frame = open(frame_basename+".dat", "w")
                file_dat_frame.write("#k_iter res_norm err_res err_res_rel relax dU_norm U_norm err_dU im_diff err_im\n")
Martin Genet's avatar
Martin Genet committed

                file_pvd_frame = dolfin.File(frame_basename+"_.pvd")
                file_pvd_frame << (U, 0.)

            print_str(tab,"Loading image, image gradient and image hessian…")
            image_filename = images_folder+"/"+images_basename+"_"+str(k_frame).zfill(images_zfill)+".vti"
            Idef.init_image(image_filename)
            DIdef.init_image(image_filename)
                DDIdef.init_image(image_filename)
            image_filename = images_folder+"/"+images_basename+"_"+str(k_frame_old).zfill(images_zfill)+".vti"
            Iold.init_image(image_filename)
            DIold.init_image(image_filename)
Martin Genet's avatar
Martin Genet committed
            # linear system: matrix
                print_str(tab,"Assembly…")
                A = dolfin.assemble((1.-regul_level) * DDpsi_c_old * dX, tensor=A, form_compiler_parameters={'quadrature_degree':images_quadrature})
                A = dolfin.assemble((   regul_level) * DDpsi_m     * dX, tensor=A, add_values=True)
Martin Genet's avatar
Martin Genet committed
                #print_var(tab,"A",A.array())
                #A_norm = A.norm("l2")
Martin Genet's avatar
Martin Genet committed
                #print_var(tab,"A_norm",A_norm)
Martin Genet's avatar
Martin Genet committed

Martin Genet's avatar
Martin Genet committed
            if (print_iterations):
                U.vector().zero()
                im_diff = (dolfin.assemble(psi_c * dX, form_compiler_parameters={'quadrature_degree':images_quadrature})/mesh_V0)**(1./2)
                err_im = im_diff/Iref_norm
                file_dat_frame.write(" ".join([str(val) for val in [-2, None, None, None, None, None, None, im_diff, err_im, None]])+"\n")
                im_diff = (dolfin.assemble(psi_c * dX, form_compiler_parameters={'quadrature_degree':images_quadrature})/mesh_V0)**(1./2)
                err_im = im_diff/Iref_norm
                file_dat_frame.write(" ".join([str(val) for val in [-1, None, None, None, None, None, None, im_diff, err_im, None]])+"\n")
Martin Genet's avatar
Martin Genet committed

Martin Genet's avatar
Martin Genet committed
            print_str(tab,"Running registration…")
            tab += 1
Martin Genet's avatar
Martin Genet committed
            k_iter = 0
            if   (residual_type.startswith("Iref")):
                using_Iold_residual = False
            elif (residual_type.startswith("Iold")):
                using_Iold_residual = True
Martin Genet's avatar
Martin Genet committed
            while (True):
Martin Genet's avatar
Martin Genet committed
                print_var(tab-1,"k_iter",k_iter)
Martin Genet's avatar
Martin Genet committed
                n_iter_tot += 1

Martin Genet's avatar
Martin Genet committed
                # linear system: matrix assembly
Martin Genet's avatar
Martin Genet committed
                    print_str(tab,"Matrix assembly…")
                    A = dolfin.assemble((1.-regul_level) * DDpsi_c * dX, tensor=A, form_compiler_parameters={'quadrature_degree':images_quadrature})
                    A = dolfin.assemble((   regul_level) * DDpsi_m * dX, tensor=A, add_values=True)
Martin Genet's avatar
Martin Genet committed
                    #print_var(tab,"A",A.array())
                    #A_norm = A.norm("l2")
Martin Genet's avatar
Martin Genet committed
                    #print_sci(tab,"A_norm",A_norm)
                # linear system: residual assembly
                if (k_iter > 0):
Martin Genet's avatar
Martin Genet committed
                    if (k_iter == 1):
                        B_old = B.copy()
                    elif (k_iter > 1):
                        B_old[:] = B[:]
                    res_old_norm = res_norm
                print_str(tab,"Residual assembly…")
Martin Genet's avatar
Martin Genet committed
                    B = dolfin.assemble(- (1.-regul_level) * Dpsi_c_old * dX, tensor=B, form_compiler_parameters={'quadrature_degree':images_quadrature})
                    B = dolfin.assemble(- (   regul_level) * Dpsi_m     * dX, tensor=B, add_values=True)
Martin Genet's avatar
Martin Genet committed
                    B = dolfin.assemble(- (1.-regul_level) * Dpsi_c * dX, tensor=B, form_compiler_parameters={'quadrature_degree':images_quadrature})
                    B = dolfin.assemble(- (   regul_level) * Dpsi_m * dX, tensor=B, add_values=True)
Martin Genet's avatar
Martin Genet committed
                #print_var(tab,"B",B.array())
                res_norm = B.norm("l2")
                err_res = res_norm/res_norm0
                print_sci(tab,"err_res",err_res)

                if (k_iter == 0):
                    err_res_rel = 1.
                else:
                    if (k_iter == 1):
                        dB = B - B_old
                    elif (k_iter > 1):
                        dB[:] = B[:] - B_old[:]
                    dres_norm = dB.norm("l2")
                    err_res_rel = dres_norm / res_old_norm
                    print_sci(tab,"err_res_rel",err_res_rel)
Martin Genet's avatar
Martin Genet committed

                # linear system: solve
Martin Genet's avatar
Martin Genet committed
                print_str(tab,"Solve…")
Martin Genet's avatar
Martin Genet committed
                dolfin.solve(A, dU.vector(), B)
Martin Genet's avatar
Martin Genet committed
                #print_var(tab,"dU",dU.vector().array())
Martin Genet's avatar
Martin Genet committed

                # relaxation
Martin Genet's avatar
Martin Genet committed
                if (relax_type == "constant"):
                    if (k_iter == 0):
                        relax = relax_init
                elif (relax_type == "aitken"):
                    if (k_iter == 0):
                        relax = relax_init
                    else:
                        relax *= (-1.) * B_old.inner(dB) / dres_norm**2
Martin Genet's avatar
Martin Genet committed
                elif (relax_type == "gss"):
                    phi = (1 + math.sqrt(5)) / 2
                    relax_a = (1-phi)/(2-phi)
                    relax_b = 1/(2-phi)
                    need_update_c = True
                    need_update_d = True
                    relax_cur = 0.
                    relax_list = []
                    relax_vals = []
Martin Genet's avatar
Martin Genet committed
                    tab += 1
Martin Genet's avatar
Martin Genet committed
                    relax_k = 0
                    while (True):
                        print_var(tab-1,"relax_k",relax_k)
                        print_sci(tab,"relax_a",relax_a)
                        print_sci(tab,"relax_b",relax_b)
                        if (need_update_c):
                            relax_c = relax_b - (relax_b - relax_a) / phi
                            relax_list.append(relax_c)
                            print_sci(tab,"relax_c",relax_c)
                            U.vector().axpy(relax_c-relax_cur, dU.vector())
                            relax_cur = relax_c
                            relax_fc  = dolfin.assemble((1.-regul_level) * psi_c * dX, form_compiler_parameters={'quadrature_degree':images_quadrature})
                            print_sci(tab,"relax_fc",relax_fc)
                            relax_fc += dolfin.assemble((   regul_level) * psi_m * dX)
                            print_sci(tab,"relax_fc",relax_fc)
                            #print_sci(tab,"relax_fc",relax_fc)
                            if (numpy.isnan(relax_fc)):
                                relax_fc = float('+inf')
                                print_sci(tab,"relax_fc",relax_fc)
                            relax_vals.append(relax_fc)
                            #print_var(tab,"relax_list",relax_list)
                            #print_var(tab,"relax_vals",relax_vals)
                        if (need_update_d):
                            relax_d = relax_a + (relax_b - relax_a) / phi
                            relax_list.append(relax_d)
                            print_sci(tab,"relax_d",relax_d)
                            U.vector().axpy(relax_d-relax_cur, dU.vector())
                            relax_cur = relax_d
                            relax_fd  = dolfin.assemble((1.-regul_level) * psi_c * dX, form_compiler_parameters={'quadrature_degree':images_quadrature})
                            print_sci(tab,"relax_fd",relax_fd)
                            relax_fd += dolfin.assemble((   regul_level) * psi_m * dX)
                            print_sci(tab,"relax_fd",relax_fd)
                            #print_sci(tab,"relax_fd",relax_fd)
                            if (numpy.isnan(relax_fd)):
                                relax_fd = float('+inf')
                                print_sci(tab,"relax_fd",relax_fd)
                            relax_vals.append(relax_fd)
                            #print_var(tab,"relax_list",relax_list)
                            #print_var(tab,"relax_vals",relax_vals)
                        if (relax_fc < relax_fd):
                            relax_b = relax_d
                            relax_d = relax_c
                            relax_fd = relax_fc
                            need_update_c = True
                            need_update_d = False
                        elif (relax_fc > relax_fd):
                            relax_a = relax_c
                            relax_c = relax_d
                            relax_fc = relax_fd
                            need_update_c = False
                            need_update_d = True
                        else: assert(0)
                        if (relax_k >= 5) and (numpy.argmin(relax_vals) > 0):
                            break
                        relax_k += 1
Martin Genet's avatar
Martin Genet committed
                    U.vector().axpy(-relax_cur, dU.vector())
                    #print_var(tab,"relax_vals",relax_vals)

                    if (print_iterations):
Martin Genet's avatar
Martin Genet committed
                        iter_basename = frame_basename+"-iter="+str(k_iter).zfill(3)
                        file_dat_iter = open(iter_basename+".dat","w")
                        file_dat_iter.write("\n".join([" ".join([str(val) for val in [relax_list[relax_k], relax_vals[relax_k]]]) for relax_k in xrange(len(relax_list))]))
                        file_dat_iter.close()
Martin Genet's avatar
Martin Genet committed
                        os.system("gnuplot -e \"set terminal pdf; set output '"+iter_basename+".pdf'; plot '"+iter_basename+".dat' u 1:2 w p title 'psi_int'\"")

                    relax = relax_list[numpy.argmin(relax_vals)]
Martin Genet's avatar
Martin Genet committed
                else:
Martin Genet's avatar
Martin Genet committed
                    assert (0), "relax_type must be \"constant\", \"aitken\" or \"gss\". Aborting."
Martin Genet's avatar
Martin Genet committed

                # solution update
                U.vector().axpy(relax, dU.vector())
                U_norm = U.vector().norm("l2")
Martin Genet's avatar
Martin Genet committed

                if (print_iterations):
Martin Genet's avatar
Martin Genet committed
                    #print_var(tab,"U",U.vector().array())
Martin Genet's avatar
Martin Genet committed
                    file_pvd_frame << (U, float(k_iter+1))
Martin Genet's avatar
Martin Genet committed

                # displacement error
Martin Genet's avatar
Martin Genet committed
                dU_norm = abs(relax) * dU.vector().norm("l2")
                if (dU_norm == 0.) and (Uold_norm == 0.) and (U_norm == 0.):
                elif (Uold_norm == 0.):
                    err_dU = dU_norm/Uold_norm
                print_sci(tab,"err_dU",err_dU)
Martin Genet's avatar
Martin Genet committed

                # image error
Martin Genet's avatar
Martin Genet committed
                if (k_iter > 0):
                im_diff = (dolfin.assemble(psi_c * dX, form_compiler_parameters={'quadrature_degree':images_quadrature})/mesh_V0)**(1./2)
                err_im = im_diff/Iref_norm
                print_sci(tab,"err_im",err_im)
Martin Genet's avatar
Martin Genet committed

                if (print_iterations):
                    file_dat_frame.write(" ".join([str(val) for val in [k_iter, res_norm, err_res, err_res_rel, relax, dU_norm, U_norm, err_dU, im_diff, err_im]])+"\n")
Martin Genet's avatar
Martin Genet committed

                # exit test
                if (tol_res is not None) and (err_res > tol_res):
                    success = False
                if (tol_res_rel is not None) and (err_res_rel > tol_res_rel):
                if (tol_dU is not None) and (err_dU > tol_dU):
                if (tol_im is not None) and (err_im > tol_im):
Martin Genet's avatar
Martin Genet committed
                if (success):
Martin Genet's avatar
Martin Genet committed
                    print_str(tab,"Nonlinear solver converged…")
Martin Genet's avatar
Martin Genet committed
                    break

Martin Genet's avatar
Martin Genet committed
                if (k_iter == n_iter_max-1):
                    if (residual_type=="Iref-then-Iold") and not (using_Iold_residual):
                        print_str(tab,"Warning! Nonlinear solver failed to converge…using Iold instead of Iref. (k_frame = "+str(k_frame)+")")
                        using_Iold_residual = True
                        U.vector()[:] = Uold.vector()[:]
                        U_norm = Uold_norm
                        k_iter = 0
                        continue
                    else:
                        print_str(tab,"Warning! Nonlinear solver failed to converge… (k_frame = "+str(k_frame)+")")
                        global_success = False
                        break
Martin Genet's avatar
Martin Genet committed

                # increment counter
                k_iter += 1

Martin Genet's avatar
Martin Genet committed
            tab -= 1

Martin Genet's avatar
Martin Genet committed
            if (print_iterations):
Martin Genet's avatar
Martin Genet committed
                file_dat_frame.close()
                os.system("gnuplot -e \"set terminal pdf; set output '"+frame_basename+".pdf'; set key box textcolor variable; set grid; set logscale y; set yrange [1e-3:1e0]; plot '"+frame_basename+".dat' u 1:3 pt 1 lw 3 title 'err_res', '' u 1:7 pt 1 lw 3 title 'err_dU', '' using 1:9 pt 1 lw 3 title 'err_im', "+str(tol_res or tol_dU or tol_im)+" lt -1 notitle; unset logscale y; set yrange [*:*]; plot '' u 1:4 pt 1 lw 3 title 'relax'\"")
Martin Genet's avatar
Martin Genet committed

            if not (success) and not (continue_after_fail):
                break

            Uold.vector()[:] = U.vector()[:]
            Uold_norm = U_norm
Martin Genet's avatar
Martin Genet committed
            mesh_V = dolfin.assemble(J*dX)
            print_sci(tab+1,"mesh_V",mesh_V)
            file_volume.write(" ".join([str(val) for val in [k_frame, mesh_V]])+"\n")

Martin Genet's avatar
Martin Genet committed
            print_str(tab,"Printing solution…")
Martin Genet's avatar
Martin Genet committed
            file_pvd << (U, float(k_frame))
Martin Genet's avatar
Martin Genet committed

            if (images_dynamic_scaling):
                p = numpy.empty((2,2))
                q = numpy.empty(2)
                p[0,0] = dolfin.assemble(Idef**2 * dX, form_compiler_parameters={'quadrature_degree':images_quadrature})
                p[0,1] = dolfin.assemble(Idef * dX, form_compiler_parameters={'quadrature_degree':images_quadrature})
                p[1,0] = p[0,1]
                p[1,1] = 1.
                q[0] = dolfin.assemble(Idef*Iref * dX, form_compiler_parameters={'quadrature_degree':images_quadrature})
                q[1] = dolfin.assemble(Iref * dX, form_compiler_parameters={'quadrature_degree':images_quadrature})
                scaling[:] = numpy.linalg.solve(p,q)
                print_var(tab,"scaling", scaling)

                im_diff = (dolfin.assemble(psi_c * dX, form_compiler_parameters={'quadrature_degree':images_quadrature})/mesh_V0)**(1./2)
                err_im = im_diff/Iref_norm
                print_sci(tab,"err_im",err_im)
Martin Genet's avatar
Martin Genet committed
        tab -= 1

Martin Genet's avatar
Martin Genet committed
        if not (success) and not (continue_after_fail):
            break

Martin Genet's avatar
Martin Genet committed
    print_var(tab,"n_iter_tot",n_iter_tot)
Martin Genet's avatar
Martin Genet committed

Martin Genet's avatar
Martin Genet committed
    file_volume.close()
Martin Genet's avatar
Martin Genet committed

Martin Genet's avatar
Martin Genet committed
    return global_success