compute_strains.py

#coding=utf8

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

import dolfin
import glob
import numpy

import myPythonLibrary as mypy
import myVTKPythonLibrary as myvtk

import dolfin_dic as ddic

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

def compute_strains(
        sol_folder,
        sol_basename,
        sol_ext="vtu",
        ref_frame=None,
        disp_array_name="displacement",
        defo_grad_array_name="DeformationGradient",
        strain_array_name="Strain",
        mesh_w_local_basis_folder=None,
        mesh_w_local_basis_basename=None,
        CYL_or_PPS="PPS",
        write_strains=1,
        plot_strains=1,
        write_strain_vs_radius=0,
        verbose=1):

    if (mesh_w_local_basis_folder is not None) and (mesh_w_local_basis_basename is not None):
        mesh_w_local_basis_filename = mesh_w_local_basis_folder+"/"+mesh_w_local_basis_basename+"-WithLocalBasis.vtk"
        mesh_w_local_basis = myvtk.readUGrid(
            filename=mesh_w_local_basis_filename,
            verbose=verbose)
        mesh_w_local_basis_n_cells = mesh_w_local_basis.GetNumberOfCells()
        if (verbose): print "mesh_w_local_basis_n_cells = " + str(mesh_w_local_basis_n_cells)

        if (mesh_w_local_basis.GetCellData().HasArray("sector_id")):
            iarray_sector_id = mesh_w_local_basis.GetCellData().GetArray("sector_id")
            n_sector_ids = 0
            for k_cell in xrange(mesh_w_local_basis_n_cells):
                sector_id = int(iarray_sector_id.GetTuple1(k_cell))
                if (sector_id < 0): continue
                if (sector_id > n_sector_ids-1):
                    n_sector_ids = sector_id+1
            if (verbose): print "n_sector_ids = " + str(n_sector_ids)
        else:
            n_sector_ids = 0
    else:
        mesh_w_local_basis = None
        n_sector_ids = 0

    sol_filenames = glob.glob(sol_folder+"/"+sol_basename+"_[0-9]*."+sol_ext)

    sol_zfill = len(sol_filenames[0].rsplit("_",1)[-1].split(".")[0])
    if (verbose): print "sol_zfill = " + str(sol_zfill)

    n_frames = len(sol_filenames)
    if (verbose): print "n_frames = " + str(n_frames)

    if (write_strains):
        strain_file = open(sol_folder+"/"+sol_basename+"-strains.dat", "w")
        strain_file.write("#t Err_avg Err_std Ecc_avg Ecc_std Ell_avg Ell_std Erc_avg Erc_std Erl_avg Erl_std Ecl_avg Ecl_std\n")

    if (write_strain_vs_radius):
        strain_vs_radius_file = open(sol_folder+"/"+sol_basename+"-strains_vs_radius.dat", "w")
        strain_vs_radius_file.write("#t r Err Ecc Ell Erc Erl Ecl\n")

    if (ref_frame is not None):
        ref_mesh_filename = sol_folder+"/"+sol_basename+"_"+str(ref_frame).zfill(sol_zfill)+"."+sol_ext
        ref_mesh = myvtk.readUGrid(
            filename=ref_mesh_filename,
            verbose=verbose)
        myvtk.addDeformationGradients(
            mesh=ref_mesh,
            disp_array_name=disp_array_name,
            verbose=verbose)
        farray_F0 = ref_mesh.GetCellData().GetArray(defo_grad_array_name)

    for k_frame in xrange(n_frames):
        mesh_filename = sol_folder+"/"+sol_basename+"_"+str(k_frame).zfill(sol_zfill)+"."+sol_ext
        mesh = myvtk.readUGrid(
            filename=mesh_filename,
            verbose=verbose)
        n_cells = mesh.GetNumberOfCells()
        if (mesh_w_local_basis is not None):
            assert (n_cells == mesh_w_local_basis_n_cells)
            mesh.GetCellData().AddArray(mesh_w_local_basis.GetCellData().GetArray("part_id"))
            mesh.GetCellData().AddArray(mesh_w_local_basis.GetCellData().GetArray("sector_id"))
        myvtk.addDeformationGradients(
            mesh=mesh,
            disp_array_name=disp_array_name,
            defo_grad_array_name=defo_grad_array_name,
            verbose=verbose)
        if (ref_frame is not None):
            farray_F = mesh.GetCellData().GetArray(defo_grad_array_name)
            for k_cell in xrange(n_cells):
                F  = numpy.reshape(farray_F.GetTuple(k_cell) , (3,3), order='C')
                F0 = numpy.reshape(farray_F0.GetTuple(k_cell), (3,3), order='C')
                F  = numpy.dot(F, numpy.linalg.inv(F0))
                farray_F.SetTuple(k_cell, numpy.reshape(F, 9, order='C'))
        myvtk.addStrainsFromDeformationGradients(
            mesh=mesh,
            defo_grad_array_name=defo_grad_array_name,
            strain_array_name=strain_array_name,
            mesh_w_local_basis=mesh_w_local_basis,
            verbose=verbose)
        myvtk.writeUGrid(
            ugrid=mesh,
            filename=sol_folder+"/"+sol_basename+"_"+str(k_frame).zfill(sol_zfill)+"."+sol_ext,
            verbose=verbose)

        if (write_strains) or (write_strain_vs_radius):
            if (mesh_w_local_basis is not None):
                assert (mesh.GetCellData().HasArray(strain_array_name+"_"+CYL_or_PPS))
                farray_strain = mesh.GetCellData().GetArray(strain_array_name+"_"+CYL_or_PPS)
            else:
                farray_strain = mesh.GetCellData().GetArray(strain_array_name)

        if (write_strains):
            if (n_sector_ids == 0):
                strains_all = []
                for k_cell in xrange(n_cells):
                    strains_all.append(farray_strain.GetTuple(k_cell))
            elif (n_sector_ids == 1):
                strains_all = []
                for k_cell in xrange(n_cells):
                    sector_id = int(iarray_sector_id.GetTuple(k_cell)[0])
                    if (sector_id < 0): continue
                    strains_all.append(farray_strain.GetTuple(k_cell))
            elif (n_sector_ids > 1):
                strains_all = []
                strains_per_sector = [[] for sector_id in xrange(n_sector_ids)]
                for k_cell in xrange(n_cells):
                    sector_id = int(iarray_sector_id.GetTuple(k_cell)[0])
                    if (sector_id < 0): continue
                    strains_all.append(farray_strain.GetTuple(k_cell))
                    strains_per_sector[sector_id].append(farray_strain.GetTuple(k_cell))

            strain_file.write(str(k_frame))
            strains_all_avg = numpy.mean(strains_all, 0)
            strains_all_std = numpy.std(strains_all, 0)
            strain_file.write("".join([" " + str(strains_all_avg[k_comp]) + " " + str(strains_all_std[k_comp]) for k_comp in xrange(6)]))
            if (n_sector_ids > 1):
                for sector_id in xrange(n_sector_ids):
                    strains_per_sector_avg = numpy.mean(strains_per_sector[sector_id], 0)
                    strains_per_sector_std = numpy.std(strains_per_sector[sector_id], 0)
                    strain_file.write("".join([" " + str(strains_per_sector_avg[k_comp]) + " " + str(strains_per_sector_std[k_comp]) for k_comp in xrange(6)]))
            strain_file.write("\n")

        if (write_strain_vs_radius):
            assert (mesh_w_local_basis.GetCellData().HasArray("r"))
            farray_r = mesh_w_local_basis.GetCellData().GetArray("r")
            for k_cell in xrange(n_cells):
                strain_vs_radius_file.write(" ".join([str(val) for val in [k_frame, farray_r.GetTuple1(k_cell)]+list(farray_strain.GetTuple(k_cell))]) + "\n")
            strain_vs_radius_file.write("\n")
            strain_vs_radius_file.write("\n")

    if (write_strains):
        strain_file.close()

        if (plot_strains):
            ddic.plot_strains(
                working_folder=sol_folder,
                working_basenames=[sol_basename],
                suffix=None,
                verbose=verbose)

    if (write_strain_vs_radius):
        strain_vs_radius_file.close()