stpWriter.py

'''
stpWriter
'''

import sys, os
import argparse

from OCC.gp import gp_Pnt, gp_Dir, gp_Vec, gp_Pln, gp_Ax3
from OCC.Utils.DataExchange.STEP import STEPExporter
from OCC.TopoDS import TopoDS_Builder
from OCC.TopoDS import TopoDS_Compound
from OCC.TNaming import TNaming_Builder
from OCC.TDF import TDF_Label
import OCC.TDocStd as TDocStd
import OCC.TCollection as TCollection
import OCC.AppStd as AppStd
from OCC.BRepBuilderAPI import BRepBuilderAPI_MakeWire, BRepBuilderAPI_MakeEdge
from OCC.TopoDS import TopoDS_edge, TopoDS_vertex, TopoDS_wire

title = 'stpWriter'

def read_keyword_def(lines, dico_output):

    '''
        Read the parameters corresponding to a keyword defined in a BDF file
        (possibly on several lines).
    '''

    linesToMod = lines[:]
    drapAng    = None
    try:
        # Combine lines with *
        star_lines = []
        for k in range(len(linesToMod)):
            if '*' in linesToMod[k][1:]:
                star_lines.append(k)
        star_lines.sort()
        star_lines.reverse()
        for Lind in star_lines:
            temp_line = linesToMod[Lind].strip().ljust(72)[:72]
            # should the ending be used?
            ending = linesToMod[Lind].strip().ljust(72)[72:]
            try:
                if linesToMod[Lind+1][0] == '*':
                    temp_line += linesToMod[Lind+1][8:].rstrip().ljust(64)
                    linesToMod[Lind] = temp_line
                    del linesToMod[Lind+1]
                else:
                    continue
            except IndexError:
                pass

        # Read the first line
        line = linesToMod[0].strip()
        if '*' in line: lcd = 16
        else:           lcd = 8
        typ  = line[:8].replace('*','').strip().upper()
        name = line[8:8+lcd].strip().upper()
        data = []
        parameters = []

        for k in range(8+lcd,8+lcd*9,lcd):
            if k >= len(line): parameters.append('')
            else:              parameters.append(line[k:k+lcd].strip())
        data.append(parameters)

        # Read the following lines for the same keyword
        for line in linesToMod[1:]:
            parameters = []
            line       = line.replace('\n','').replace('\r','')

            # For a PCOMP, record in a list the first line and
            # then each line corresponding to a ply definition
            if typ == 'PCOMP':
                if 'draping angle' in line.lower():
                    line.replace(':', '=')
                    drapAng = str(line.split('=')[-1].replace('deg', '').strip())
                    continue

                for k in range(0,4*lcd,lcd):
                    if k+8 >= len(line)*lcd/8:
                        parameters.append('')
                    else:
                        parameters.append(line[k+8:k+8+lcd].strip())
                data.append(parameters)
                #explain why use just one field in the middle
                if line[5*lcd:5*lcd+lcd].strip():
                    parameters = []
                    for k in range(5*lcd,9*lcd,lcd):
                        if k >= len(line)*lcd/8:
                            parameters.append('')
                        else:
                            parameters.append(line[k:k+lcd].strip())
                    data.append(parameters)

            elif typ == 'CORD2R':
                for k in range(8,len(line),lcd):
                    parameters.append(line[k:k+lcd].strip())
                data.append(parameters)

            # For a PSHELL or MATi, record in a list each line
            # corresponding to its definition
            else:
                for k in range(8,9*lcd,lcd):
                    if k >= len(line)*lcd/8:
                        parameters.append('')
                    else:
                        parameters.append(line[k:k+lcd].strip())
                data.append(parameters)

        # Add inputs to dictionary and return
        if typ not in ['FORCE','MOMENT']:
            dico_output.setdefault(typ, {})[name] = data
        else:
        # for FORCE and MOMENT another hierarchy-level is needed
            gridID = data[0][0]
            if gridID in dico_output.get(typ, {}).get(name,{}):
                olddata = dico_output[typ][name].pop(gridID)
                for i in range(3,6):
                    oldValue = conv2float(olddata[0][i])
                    toAdd = conv2float(data[0][i])
                    merged = oldValue + toAdd
                    data[0][i] = normalize_Float(merged, 15).replace(' ','')
            dico_output.setdefault(typ, {}).setdefault(name,{})[gridID] = data

        if drapAng: dico_output.setdefault('DRAPING ANGLE', {})[name] = drapAng
    except Exception as currentEx:
        print('An exception happened during reading a set of lines: ' + str(lines))

    return dico_output


def read_bdf(bdf_path=None, default_path='*.bdf', keywords = [],
             assume_bulk = 0, skip = 0, flagREST = 0):

    '''
        Reads the BDF input file given in parameter (pops up a window to select a
        file when "None" is used) and returns a dictionary containing the following
        objects encountered:

            MAT1, MAT8, PCOMP, PSHELL, GRID, CTRIA3, CQUAD4, CBUSH, SUBCASE

        dico_FE_matprop[object_type][object_ID] = [[ Line1_parameters ],
                                                   [ Line2_parameters ] ...]

        dico_FE_matprop['SUBCASE'][subcase_ID] = [subtitle, temperature_value]

        assume_bulk: Set to 1 if you want to read in an INCLUDE file without BEGIN BULK
        skip       : Set to 1 if you want to skip all INCLUDEs in bdf/dat file
        flagREST   : Set to 1 if you want to put all lines which are not recognized
                        or not covered by keyword under the keyword 'REST'

    '''

    if not keywords:
        keywords = ['MAT1'  , 'MAT8'  , 'PCOMP', 'PSHELL', 'GRID'   ,
                    'CTRIA3', 'CQUAD4', 'CBUSH', 'CORD2R' ]

    # Target the BDF file to be read and read lines
    bdf_path   =  bdf_path.replace('\\','/')
    bdf_paths  = [bdf_path]
    Bulk_Flags = [assume_bulk]              # This flag indicates that the keyword "BEGIN BULK"
                                              # has not yet been read for the BDF file in "bdf_paths"

    # Read shell and materials definition
    dico_FE_matprop = {}
    subtitle    = 'NC'
    temp_id     = 'NC'
    temp_id_flag    = 0

    temp_def    = {'NC': 'NC'}

    for bdp in bdf_paths:

        flag_Bulk = Bulk_Flags[bdf_paths.index(bdp)]

        try: bdfi = open(bdp)
        except:
            incl_name = bdp.split('/')[-1]
            incl_path = eg.fileopenbox('Pick up the BDF include file: "'+
                                       incl_name+'"', '', default_path)
            if incl_path == None:
                bdfi = ''
            else:
                bdfi = open(incl_path,'rb')

        if bdfi != '':

            data_lines = []
            subcase_id = ''

            i = 0
            for line in bdfi:

                i += 1

                if not flag_Bulk:

                    typ = line.strip().upper()

                    if typ[:7] == 'SUBCASE':

                        if subcase_id != '':
                            dico_FE_matprop.setdefault('SUBCASE',{}).setdefault(subcase_id,[subtitle, temp_id])

                        subcase_id   = reanyinteger.findall(line)[0]
                        subtitle     = 'NC'
                        temp_id   = 'NC'
                        temp_id_flag = 0

                    elif typ[:8] == 'SUBTITLE':
                        subtitle = line.replace('SUBTITLE','').replace('=','',1).strip()

                    elif typ[:10] == 'TEMP(LOAD)' or typ[:17] == 'TEMPERATURE(LOAD)':
                        temp_id = reanyinteger.findall(line)[0]
                        temp_id_flag += 1

                    elif typ[:13] == 'TEMP(INITIAL)' or typ[:20] == 'TEMPERATURE(INITIAL)':
                        dico_FE_matprop.setdefault('SUBCASE',{})['INITIAL'] = ['INITIAL', reanyinteger.findall(line)[0]]

                    elif typ[:10] == 'BEGIN BULK':

                        if subcase_id != '':
                            dico_FE_matprop.setdefault('SUBCASE',{}).setdefault(subcase_id,[subtitle, temp_id])

                        flag_Bulk = 1

                    elif typ[:7] == 'INCLUDE' and not skip:
                        incl_name = line.split("'")[-2]
                        incl_path = '/'.join(bdf_path.split('/')[:-1]) + '/' + incl_name
                        bdf_paths.append(incl_path)
                        Bulk_Flags.append(flag_Bulk)

                else:

                    typ = line[:8].replace('*',' ').strip().upper()

                    if data_lines:
                        if line[:8].replace('*',' ') == ' '*8:
                            data_lines.append(line)
                            continue
                        elif line[0] == '$':
                            if 'draping angle' in line.lower():
                                data_lines.append(line)
                                continue
                            else:
                                continue
                        elif '*' in line and len(line) < 8:
                            data_lines.append(line)
                            continue
                        else:
                            dico_FE_matprop = read_keyword_def(data_lines, dico_FE_matprop)
                            data_lines    = []

                    if typ in keywords:
                        data_lines.append(line)

                    elif typ == 'TEMPD':
                        temp_def[reanyinteger.findall(line)[0]] = reanynumber2.findall(line)[-1]

                    elif typ == 'INCLUDE' and not skip:
                        incl_name = line.split("'")[-2]
                        incl_path = '/'.join(bdf_path.split('/')[:-1]) + '/' + incl_name
                        bdf_paths.append(incl_path)
                        Bulk_Flags.append(flag_Bulk)
                    #useful if bdf file should be read in and written out in a similar way
                    elif flagREST:
                        dico_FE_matprop.setdefault('REST', []).append((i,line))


            if data_lines:
                dico_FE_matprop = read_keyword_def(data_lines, dico_FE_matprop)
                data_lines    = []

            if subcase_id != '':
                dico_FE_matprop['SUBCASE'][subcase_id] = [subtitle, temp_id]

            bdfi.close()

    # Check if the temperatures are defined in C or K
    temp_li = []
    for i in temp_def:
        try: temp_li.append(float(temp_def[i]))
        except: None
    if temp_li != [] and max(temp_li) > 200: # It is assumed that temperatures are defined in K in the model
        for i in temp_def:
            try: temp_def[i] = str(float(temp_def[i]) - 273.15)
            except: None
    nonValids = set()
    # Check that material definitions #MAT1 are correct
    if 'MAT1' in dico_FE_matprop:
        for mid in dico_FE_matprop['MAT1']:

            missing = []
            try: E = float(dico_FE_matprop['MAT1'][mid][0][0])
            except: missing.append('E')
            try: G = float(dico_FE_matprop['MAT1'][mid][0][1])
            except: missing.append('G')
            try: v = float(dico_FE_matprop['MAT1'][mid][0][2])
            except: missing.append('v')

            if len(missing) == 1:
                if missing[0] == 'E':
                    E = 2*G*(1+v)
                    dico_FE_matprop['MAT1'][mid][0][0] = str(E)
                elif missing[0] == 'G':
                    G = E/(2*(1+v))
                    dico_FE_matprop['MAT1'][mid][0][1] = str(G)
                elif missing[0] == 'v':
                    v = E/(2*G) - 1
                    dico_FE_matprop['MAT1'][mid][0][2] = str(v)

            elif len(missing) > 1:
                nonValids.add(mid)
        while nonValids:
            nonValids.pop()
            del dico_FE_matprop['MAT1'][mid]

    # Check that material definitions #MAT8 are correct

    if 'MAT8' in dico_FE_matprop:
        for mid in dico_FE_matprop['MAT8']:

            missing = []
            try:    float(dico_FE_matprop['MAT8'][mid][0][0])
            except: missing.append('El')
            try:    float(dico_FE_matprop['MAT8'][mid][0][1])
            except: missing.append('Et')
            try:    float(dico_FE_matprop['MAT8'][mid][0][3])
            except: missing.append('Glt')
            try:    float(dico_FE_matprop['MAT8'][mid][0][2])
            except: missing.append('vlt')

            if len(missing) != 0:
                nonValids.add(mid)
        while nonValids:
            nonValids.pop()
            del dico_FE_matprop['MAT8'][mid]

    # Complete the subcases temperatures definition

    if dico_FE_matprop.get('SUBCASE',{}).get('INITIAL',['NC','NC'])[1] == 'NC':
        temp_def['DEFAULT_INIT_20'] = '20.'

    LC_without_T = []
    for subcase_id in dico_FE_matprop.get('SUBCASE',{}):
        if dico_FE_matprop['SUBCASE'][subcase_id][1] == 'NC':
            LC_without_T.append(subcase_id)
            init_id = dico_FE_matprop['SUBCASE']['INITIAL'][1]
            dico_FE_matprop['SUBCASE'][subcase_id][1] = init_id
    LC_without_T.sort()

    for subcase_id in dico_FE_matprop.get('SUBCASE',{}):
        temp_id = dico_FE_matprop['SUBCASE'][subcase_id][1]
        if temp_id in temp_def:
            temp_val = temp_def[temp_id]
            dico_FE_matprop['SUBCASE'][subcase_id].append(temp_val)

    return dico_FE_matprop, bdf_paths


def write_step(dicBdf, scale=1, title='Title', fileName=None, setElems=set()):

    '''
        write_step: write a step file which contains faces build from the elements from the dicBdf

    '''

    allTria = set(dicBdf.get('CTRIA3',{}).keys())
    allQuad = set(dicBdf.get('CQUAD4',{}).keys())
    allElems = allTria | allQuad
    if not setElems:
        validElems = allElems
    else:
        validElems = setElems & allElems

    validQuads = validElems & allQuad
    faceset_quad=WiresFromElements(dicBdf,validQuads,1)

    validTrias = validElems & allTria
    faceset_tria=WiresFromElements(dicBdf,validTrias,0)

    allFaces=faceset_quad | faceset_tria

    stepWriter=STEPExporter(fileName)

    label = TDF_Label()

    comp = TopoDS_Compound()
    compBuilder = TopoDS_Builder()
    compBuilder.MakeCompound(comp)
    for shape in allFaces:
        compBuilder.Add(comp, shape)
    stepWriter.add_shape(comp)
    stepWriter.write_file()

    return fileName

def WiresFromElements(bdf,elements,quad):

    wire_set=set([])
    for elem in elements:
        if quad:
            nodes=bdf['CQUAD4'][elem][0][1:5]
        elif not quad:
            nodes=bdf['CTRIA3'][elem][0][1:4]
        pnt_list=[]
        edge_list=[]
        for node in nodes:
            x=conv2float(bdf['GRID'][node][0][1])
            y=conv2float(bdf['GRID'][node][0][2])
            z=conv2float(bdf['GRID'][node][0][3])
            pnt_list.append(gp_Pnt(x,y,z))
        for i,point in enumerate(pnt_list):
            edge=BRepBuilderAPI_MakeEdge(pnt_list[i-1],pnt_list[i]).Edge()
            edge_list.append(edge)
        buildWire=makeWirefromWires()
        buildWire.Add(edge_list,1)
        wire=buildWire.Build()
        wire_set.add(wire)

    return wire_set

def conv2float(val):

    try:
        return float(val)
    except:
        val = val[0] + val[1:].replace('-','E-').replace('+','E+')
        return float(val)

class makeWirefromWires(object):

    def __init__(self):
        self.buildWire=BRepBuilderAPI_MakeWire()

    def Add(self,wire,flagEdgeList=0):

        edgelist=[]
        if not flagEdgeList:
            exp=TopExp_Explorer(wire,TopAbs_EDGE)
            while exp.More():
                edgelist.append(exp.Current())
                exp.Next()
        else:
            edgelist=wire

        count=0
        while edgelist:
            count+=1
            if count==1000:
                print 'CAUTION:'
                print 'When your wire consists of more than 1000 edges... Sorry!:'
                print ' -> Please change the number of counts in  class makeWirefromWires in occ_module.py'
                print 'Otherwise the function Add in class BRepBuilderAPI_MakeWire failed:'
                print ' -> please check the order of your added wires or set down the step tolerance'
                exit()

            for edge in edgelist:
                self.buildWire.Add(TopoDS_edge(edge))
                if self.buildWire.IsDone():
                    edgelist.remove(edge)

    def Build(self):
        topoWire = self.buildWire.Wire()
        return topoWire

if __name__ == '__main__':

    # command line arguments parser
    p = argparse.ArgumentParser(description='Read bdf and save as step.', epilog='Example:\n./stpWriter.py input.bdf output.step [options]')

    # postitional arguments
    p.add_argument("infile", help="Input file")
    p.add_argument("outfile", help="Output file")

    # parse everything
    args = p.parse_args()

    bdfFile = args.infile
    stepFile = args.outfile

    dicBdf, paths = read_bdf(bdf_path=bdfFile, assume_bulk=1)
    write_step(dicBdf,
              fileName = stepFile,
              setElems = set([]))

## END OF stpWriter.py