Skip to content
Extraits de code Groupes Projets
mesh_add_rivers.py 3,72 ko
Newer Older
  • Learn to ignore specific revisions
  • import gmsh
    import numpy as np
    from scipy.spatial import cKDTree
    import getopt
    import sys
    import param
    
    gmsh.initialize()
    
    opts = None
    args = None
    try:
        opts, args = getopt.getopt(sys.argv[1:], "m:")
    except getopt.GetoptError:
        print("Argument error!")
        sys.exit(1)
    
    mesh_setup = "gbr_styx"
    
    for opt, arg in opts:
        if opt == "-m":
            mesh_setup = arg
    
    p = param.parameters(mesh_setup)
    
    
    def read_rivers(fname):
        rivers = []
        with open(fname, "r") as f:
            f.readline()
            for l in f:
                w = l.split(",")
                rivers.append(((w[3].strip()[1:-1]), (float(w[0]), float(w[1]))))
        return rivers
    
    
    def get_entities_from_physical_names(pnames):
        tags = []
        for pdim, ptag in gmsh.model.get_physical_groups(1):
            pname = gmsh.model.get_physical_name(pdim, ptag)
            if pname not in pnames:
                continue
            tags.extend(gmsh.model.get_entities_for_physical_group(pdim, ptag))
        return tags
    
    
    def get_segments(entities):
        segment_tags = []
        segment_nodes = []
        for tag in entities:
            etags, enodes = gmsh.model.mesh.getElementsByType(1, tag)
            segment_tags.append(etags)
            segment_nodes.append(enodes)
    
        segment_tags = np.hstack(segment_tags)
        segment_nodes = np.hstack(segment_nodes)
        xtag, x, _ = gmsh.model.mesh.get_nodes()
        x = x.reshape([-1, 3])
        xsearch = np.argsort(xtag)
        ex = x[xsearch[np.searchsorted(xtag, segment_nodes, sorter=xsearch)]].reshape(
            -1, 2, 3
        )
        segment_x = (ex[:, 0, :2] + ex[:, 1, :2]) / 2
        return segment_tags, segment_nodes.reshape([-1, 2]), segment_x
    
    
    def new_physical_segment(element_tag, element_node, pname):
        ent = gmsh.model.add_discrete_entity(1)
        gmsh.model.mesh.add_elements_by_type(ent, 1, [element_tag], element_node)
        ptag = gmsh.model.add_physical_group(1, [ent], -1)
        gmsh.model.set_physical_name(1, ptag, pname)
    
    
    def copy_mesh_filter_segment(eset):
        m = {}
        p = {}
        for e in gmsh.model.getEntities():
            m[e] = (
                gmsh.model.getBoundary([e]),
                gmsh.model.mesh.getNodes(e[0], e[1]),
                gmsh.model.mesh.getElements(e[0], e[1]),
            )
        for dim, tag in gmsh.model.get_physical_groups():
            p[(dim, tag)] = (
                gmsh.model.get_physical_name(dim, tag),
                gmsh.model.get_entities_for_physical_group(dim, tag),
            )
        gmsh.model.add("model2")
        for e in sorted(m):
            gmsh.model.addDiscreteEntity(e[0], e[1], [b[1] for b in m[e][0]])
            gmsh.model.mesh.addNodes(e[0], e[1], m[e][1][0], m[e][1][1])
            for i in range(m[e][2][0].size):
                if m[e][2][0][i] == 1:
                    keep = np.array(list(a not in eset for a in m[e][2][1][i]))
                    m[e][2][1][i] = m[e][2][1][i][keep]
                    m[e][2][2][i] = m[e][2][2][i].reshape([-1, 2])[keep].reshape([-1])
            gmsh.model.mesh.addElements(e[0], e[1], m[e][2][0], m[e][2][1], m[e][2][2])
        for (dim, tag), (name, ents) in p.items():
            gmsh.model.add_physical_group(dim, ents, tag)
            gmsh.model.set_physical_name(dim, tag, name)
    
    
    gmsh.open(p.mesh_file)
    rivers = read_rivers(
        f"{p.local_base_dir}/spatial/geoflows_australia_on_coast_utm56.csv"
    )
    entities = get_entities_from_physical_names(p.closed_tags)
    segment_tags, segment_nodes, segment_x = get_segments(entities)
    segment_tree = cKDTree(segment_x)
    segment_rivers = {}
    for (rtag, rx) in rivers:
        sid = segment_tree.query([rx])[1][0]
        segment_rivers[rtag] = (segment_tags[sid], segment_nodes[sid])
    copy_mesh_filter_segment(set(map(lambda a: a[0], segment_rivers.values())))
    for rtag, (stag, snodes) in segment_rivers.items():
        new_physical_segment(stag, snodes, rtag)
    
    gmsh.option.set_number("Mesh.MshFileVersion", 2)
    gmsh.write(f"{p.mesh_file[:-4]}_with_rivers.msh")