JuliaFEM.jl/src/preprocess.jl

# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

#=
- read meshes from different formats
- reorder connectivity, create element sets, node sets, ...
- create partitions for parallel runs
- renumber elements / nodes
- maybe precheck for bad elements
- check surface normal direction in boundary elements
- orientation of 2d elements
- etc only topology related stuff
=#

mutable struct Mesh
    nodes::Dict{Int,Vector{Float64}}
    node_sets::Dict{Symbol,Set{Int}}
    elements::Dict{Int,Vector{Int}}
    element_types::Dict{Int,Symbol}
    element_codes::Dict{Int,Symbol}
    element_sets::Dict{Symbol,Set{Int}}
    surface_sets::Dict{Symbol,Vector{Tuple{Int,Symbol}}}
    surface_types::Dict{Symbol,Symbol}
end

function Mesh()
    return Mesh(Dict(), Dict(), Dict(), Dict(), Dict(), Dict(), Dict(), Dict())
end

"""
    Mesh(m::Dict)

Create a new `Mesh` using data `m`. It is assumed that `m` is in format what
`abaqus_read_mesh` in `AbaqusReader.jl` is returning.
"""
function Mesh(m::Dict)
    mesh = Mesh()
    mesh.nodes = m["nodes"]
    mesh.elements = m["elements"]
    mesh.element_types = m["element_types"]
    for (k, v) in m["surface_types"]
        mesh.surface_types[Symbol(k)] = v
    end
    for (nset_name, node_ids) in m["node_sets"]
        mesh.node_sets[Symbol(nset_name)] = Set(node_ids)
    end
    for (elset_name, element_ids) in m["element_sets"]
        mesh.element_sets[Symbol(elset_name)] = Set(element_ids)
    end
    for (surfset_name, surfaces) in m["surface_sets"]
        mesh.surface_sets[Symbol(surfset_name)] = surfaces
    end
    return mesh
end

"""
    add_node!(mesh, nid, ncoords)

Add node into the mesh. `nid` is node id and `ncoords` are the node
coordinates.
"""
function add_node!(mesh::Mesh, nid::Int, ncoords::Vector{Float64})
    mesh.nodes[nid] = ncoords
end

"""
    add_nodes!(mesh, nodes)

Add nodes into the mesh.
"""
function add_nodes!(mesh::Mesh, nodes::Dict{Int,Vector{Float64}})
    for (nid, ncoords) in nodes
        add_node!(mesh, nid, ncoords)
    end
end

"""
    add_node_to_node_set!(mesh, nid, ncoords)

Add nodes into a node set. `set_name` is the name of the set and `nids...`
are all the node id:s that wants to be added.
"""
function add_node_to_node_set!(mesh::Mesh, set_name, nids...)
    if !haskey(mesh.node_sets, set_name)
        mesh.node_sets[set_name] = Set{Int}()
    end
    push!(mesh.node_sets[set_name], nids...)
    return
end

"""
    create_node_set_from_element_set!(mesh, set_names...)

Create a new node set from the nodes in an element set. ´set_names...´ are all
the set names to be inserted in the function.
"""
function create_node_set_from_element_set!(mesh::Mesh, set_names::String...)
    for set_name in set_names
        set_name = Symbol(set_name)
        @info("Creating node set $set_name from element set")
        node_ids = Set{Int}()
        for elid in mesh.element_sets[set_name]
            push!(node_ids, mesh.elements[elid]...)
        end
        mesh.node_sets[set_name] = node_ids
    end
    return
end

"""
    create_node_set_from_element_set!(mesh, set_name)

Create a new node set from an element set.
"""
function create_node_set_from_element_set!(mesh::Mesh, set_name::Symbol)
    create_node_set_from_element_set!(mesh, string(set_name))
end

"""
    add_element!(mesh, elid, eltype, connectivity)

Add an element into the mesh. ´elid´ is the element id, ´eltype´ is the type of
the element and ´connectivity´ is the connectivity of the element.
"""
function add_element!(mesh::Mesh, elid, eltype, connectivity)
    mesh.elements[elid] = connectivity
    mesh.element_types[elid] = eltype
    return nothing
end

"""
    add_elements!(mesh, elements)

Add elements into the mesh.
"""
function add_elements!(mesh::Mesh, elements::Dict{Int,Tuple{Symbol,Vector{Int}}})
    for (elid, (eltype, elcon)) in elements
        add_element!(mesh, elid, eltype, elcon)
    end
    return nothing
end

"""
    add_element_to_element_set!(mesh, set_name, elids...)

Add elements into the mesh. ´set_name´ is the name of the element set and
´elids..´ are id:s of all the elements that wants to be added.
"""
function add_element_to_element_set!(mesh::Mesh, set_name, elids...)
    if !haskey(mesh.element_sets, set_name)
        mesh.element_sets[set_name] = Set{Int}()
    end
    push!(mesh.element_sets[set_name], elids...)
end

"""
    copy(mesh)

Return a copy of the mesh.
"""
function Base.copy(mesh::Mesh)
    mesh2 = Mesh()
    mesh2.nodes = copy(mesh.nodes)
    mesh2.node_sets = copy(mesh.node_sets)
    mesh2.elements = copy(mesh.elements)
    mesh2.element_types = copy(mesh.element_types)
    mesh2.element_sets = copy(mesh.element_sets)
    return mesh2
end

"""
    filter_by_element_id(mesh, element_ids)

Filter elements by their id's.
"""
function filter_by_element_id(mesh::Mesh, element_ids::Vector{Int})
    mesh2 = copy(mesh)
    mesh2.elements = Dict()
    for elid in element_ids
        if haskey(mesh.elements, elid)
            mesh2.elements[elid] = mesh.elements[elid]
        end
    end
    return mesh2
end

"""
    filter_by_element_set(mesh, set_name)

Filter elements by an element set.
"""
function filter_by_element_set(mesh::Mesh, set_name)
    filter_by_element_id(mesh::Mesh, collect(mesh.element_sets[set_name]))
end

"""
    create_element(mesh, id)

Create an element from the mesh by it's id.
"""
function create_element(mesh::Mesh, id::Int)
    connectivity = mesh.elements[id]
    element_type = getfield(JuliaFEM, mesh.element_types[id])
    element = Element(element_type, connectivity)
    element.id = id
    update!(element, "geometry", mesh.nodes)
    return element
end

function create_elements(mesh::Mesh; element_type=nothing)
    element_ids = collect(keys(mesh.elements))
    if element_type != nothing
        filter!(id -> mesh.element_types[id] == element_type, element_ids)
    end
    elements = [create_element(mesh, id) for id in element_ids]
    return elements
end

function create_elements(mesh::Mesh, element_sets::Symbol...; element_type=nothing)
    if isempty(element_sets)
        element_ids = collect(keys(mesh.elements))
    else
        element_ids = Set{Int}()
        for set_name in element_sets
            element_ids = union(element_ids, mesh.element_sets[set_name])
        end
    end

    if element_type != nothing
        filter!(id -> mesh.element_types[id] == element_type, element_ids)
    end

    elements = [create_element(mesh, id) for id in element_ids]

    nelements = length(elements)
    content = Dict{Symbol,Int}()
    for elid in element_ids
        eltype = mesh.element_types[elid]
        content[eltype] = get(content, eltype, 0) + 1
    end
    s = join(("$v x $k" for (k, v) in content), ", ")
    v = join(element_sets, ", ")
    @info("Created $nelements elements ($s) from element set: $v.")

    return elements
end

"""
    create_elements(mesh::Mesh, element_set::String)

# Examples

Suppose that there is a `mesh` with element set `Body_1`. Creating elements
based on that element set is done then

```julia
create_elements(mesh, "Body_1")
```
"""
function create_elements(mesh::Mesh, element_sets::String...)
    return create_elements(mesh, map(Symbol, element_sets)...)
end


"""
    find_nearest_nodes(mesh, coords, npts=1; node_set=nothing)

find npts nearest nodes from the mesh and return their id numbers as a list.
"""
function find_nearest_nodes(mesh::Mesh, coords::Vector{Float64}, npts::Int=1; node_set=nothing)
    dist = Dict{Int,Float64}()
    for (nid, c) in mesh.nodes
        if node_set != nothing && !(nid in mesh.node_sets[Symbol(node_set)])
            continue
        end
        dist[nid] = norm(coords - c)
    end
    s = sort(collect(dist), by=x -> x[2])
    nd = s[1:npts] # [(id1, dist1), (id2, dist2), ..., (id_npts, dist_npts)]
    node_ids = [n[1] for n in nd]
    return node_ids
end

function find_nearest_node(mesh::Mesh, coords::Vector{Float64}; node_set=nothing)
    return first(find_nearest_nodes(mesh, coords, 1; node_set=node_set))
end

"""
    reorder_element_connectivity!(mesh, mapping)

Apply a new node ordering to elements. JuliaFEM uses the same node ordering as
ABAQUS. If the mesh is parsed from FEM format with some other node ordering,
this function can be used to reorder the nodes.

Parameters
----------
mapping :: Dict{Symbol, Vector{Int}}
    e.g. :Tet10, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

"""
function reorder_element_connectivity!(mesh::Mesh, mapping::Dict{Symbol,Vector{Int}})
    for (elid, eltype) in mesh.element_types
        haskey(mapping, eltype) || continue
        new_order = mapping[eltype]
        element_connectivity = mesh.elements[elid]
        new_element_connectivity = element_connectivity[new_order]
        mesh.elements[elid] = new_element_connectivity
    end
end

function JuliaFEM.Problem(mesh::Mesh, ::Type{P}, name::AbstractString, dimension::Int) where P<:FieldProblem
    problem = Problem(P, name, dimension)
    problem.elements = create_elements(mesh, name)
    return problem
end

function JuliaFEM.Problem(mesh::Mesh, ::Type{P}, name, dimension, parent_field_name) where P<:BoundaryProblem
    problem = Problem(P, name, dimension, parent_field_name)
    problem.elements = create_elements(mesh, name)
    return problem
end

"""
    create_coloring!(mesh::Mesh) -> Dict{Int, Int}

Greedy algorithm for coloring a grid such that no two cells with the same node
have the same color.
The returned value is a mapping between an element id and its color.
It is safe to assemble elements with the same color in parallel
"""
function create_coloring(mesh::Mesh)
    # Contains the elements that each node contain
    cell_containing_node = Dict{Int,Set{Int}}()
    for (cellid, nodes) in mesh.elements
        for v in nodes
            if !haskey(cell_containing_node, v)
                cell_containing_node[v] = Set{Int}()
            end
            push!(cell_containing_node[v], cellid)
        end
    end

    I, J, V = Int[], Int[], Bool[]
    for (node, cells) in cell_containing_node
        for cell1 in cells # All these cells have a neighboring node
            for cell2 in cells
                if cell1 != cell2
                    push!(I, cell1)
                    push!(J, cell2)
                    push!(V, true)
                end
            end
        end
    end

    incidence_matrix = sparse(I, J, V)
    # cell -> color of cell
    cell_colors = Dict{Int,Int}()
    # color -> list of cells
    final_colors = Set{Int}[]
    occupied_colors = Set{Int}()
    # Zero represents no color set yet
    for (cellid, _) in mesh.elements
        cell_colors[cellid] = 0
    end
    total_colors = 0
    for (cellid, _) in mesh.elements
        empty!(occupied_colors)
        # loop over neighbors
        for r in nzrange(incidence_matrix, cellid)
            cell_neighbour = incidence_matrix.rowval[r]
            color = cell_colors[cell_neighbour]
            if color != 0
                push!(occupied_colors, color)
            end
        end

        # occupied colors now contains all the colors we are not allowed to use
        free_color = 0
        for attempt_color in 1:total_colors
            if attempt_color ∉ occupied_colors
                free_color = attempt_color
                break
            end
        end

        if free_color == 0 # no free color found, need to bump max colors
            total_colors += 1
            free_color = total_colors
            push!(final_colors, Set{Int}())
        end

        cell_colors[cellid] = free_color
        push!(final_colors[free_color], cellid)
    end

    return cell_colors
end