fem.um

// 2D finite element method solver

import (
    "th.um"

    "matrix.um"
)

type (
    Material* = struct {
        young, poisson, yield, thickness: real
    }

    Element* = struct {
        tri: [3]int
        geometry, elasticity, stiffness: matrix::Matrix
    }
)

fn (mat: ^Material) isValid*(): bool {
    return mat.young > 0 && mat.poisson > 0 && mat.poisson < 0.5 && mat.thickness > 0
}

fn createElement*(pts: []th::Vf2, tri: [3]int, material: Material): Element {
    el := Element{tri: tri}

    a := pts[tri[1]].sub(pts[tri[0]])
    b := pts[tri[2]].sub(pts[tri[1]])
    c := pts[tri[0]].sub(pts[tri[2]])

    det := b.x * c.y - c.x * b.y
    area := 0.5 * fabs(det)

    el.geometry = matrix::Matrix{
        {-b.y,    0, -c.y,    0, -a.y,    0},
        {   0,  b.x,    0,  c.x,    0,  a.x},
        { b.x, -b.y,  c.x, -c.y,  a.x, -a.y}
    }.mulf(1.0 / det)

    el.elasticity = matrix::Matrix{
        {               1, material.poisson,                          0},
        {material.poisson,                1,                          0},
        {               0,                0, (1 - material.poisson) / 2}
    }.mulf(material.young / (1 - material.poisson * material.poisson))

    el.stiffness = el.geometry.transpose().mul(el.elasticity).mul(el.geometry).mulf(material.thickness * area)
    return el
}

fn (el: ^Element) convertDofToGlobal(dof: int): int {
    return 2 * el.tri[dof / 2] + (dof % 2)
}

fn (el: ^Element) getEquivalentStress(displacements: []th::Vf2): real {
    dofDisplacements := matrix::Matrix{
        {displacements[el.tri[0]].x},
        {displacements[el.tri[0]].y},
        {displacements[el.tri[1]].x},
        {displacements[el.tri[1]].y},
        {displacements[el.tri[2]].x},
        {displacements[el.tri[2]].y}
    }

    stresses := el.elasticity.mul(el.geometry).mul(dofDisplacements)    // {X normal stress, Y normal stress, shear stress}

    sqr := fn (x: real): real {return x * x}

    // Von Mises stress
    return sqrt(sqr(stresses[0][0]) - stresses[0][0] * stresses[1][0] + sqr(stresses[1][0]) + 3 * sqr(stresses[2][0]))
}

fn getGlobalStiffness(els: []Element, numPts: int): matrix::Matrix {
    stiffness := matrix::zeros(2 * numPts, 2 * numPts)

    for _, el in els {
        for row in el.stiffness {
            globalRow := el.convertDofToGlobal(row)
            for col in el.stiffness[0] {
                globalCol := el.convertDofToGlobal(col)
                stiffness[globalRow][globalCol] += el.stiffness[row][col]
            }
        }
    }

    return stiffness
}

fn getDofForces(forces: map[int]th::Vf2, numPts: int): matrix::Matrix {
    // Flatten force vectors into a single column vector
    dofForces := matrix::zeros(2 * numPts, 1)

    for pt, force in forces {
        if force.mag() == 0 {continue}

        dofForces[2 * pt][0] = force.x
        dofForces[2 * pt + 1][0] = force.y
    }

    return dofForces
}

fn getDofConstraints(constraints: map[int]bool): map[int]bool {
    // Each point constraint constrains both X and Y displacements
    dofConstraints := map[int]bool{}
    
    for pt, exists in constraints {
        if !exists {continue}

        dofConstraints[2 * pt] = true
        dofConstraints[2 * pt + 1] = true
    }

    return dofConstraints     
}

fn getPtDisplacements(unconstrainedDofDisplacements: matrix::Matrix, dofConstraints: map[int]bool, numPts: int): []th::Vf2 {
    displacements := make([]th::Vf2, numPts)

    dof := 0
    for pt, displacement^ in displacements {
        if dofConstraints[2 * pt] {continue}

        displacement.x = unconstrainedDofDisplacements[dof][0]
        displacement.y = unconstrainedDofDisplacements[dof + 1][0] 

        dof += 2
    }

    return displacements
}

fn getEquivalentStresses(els: []Element, displacements: []th::Vf2): []real {
    stresses := make([]real, len(els))
    for i, el in els {
        stresses[i] = el.getEquivalentStress(displacements)
    }
    return stresses
}

fn solve*(pts: []th::Vf2, tris: [][3]int, forces: map[int]th::Vf2, constraints: map[int]bool, material: Material): ([]th::Vf2, []real) {
    // The DOFs are point displacements: {displacements[0].x, displacements[0].y, displacements[1].x, displacements[1].y, ...}

    // Create finite elements
    els := make([]Element, len(tris))
    for i in tris {
        els[i] = createElement(pts, tris[i], material)
    }

    // Get global stiffness matrix
    stiffness := getGlobalStiffness(els, len(pts))
    
    // Get DOF-wise forces and constraints
    dofForces := getDofForces(forces, len(pts))
    dofConstraints := getDofConstraints(constraints)

    // Remove constrained DOFs
    unconstrainedStiffness := stiffness.remove(dofConstraints, dofConstraints)
    unconstrainedDofForces := dofForces.remove(dofConstraints, {})

    // Solve the equation: unconstrainedStiffness * unconstrainedDofDisplacements = unconstrainedDofForces
    unconstrainedDofDisplacements := unconstrainedStiffness.solve(unconstrainedDofForces)

    // Get point displacements from DOF displacements
    displacements := getPtDisplacements(unconstrainedDofDisplacements, dofConstraints, len(pts))

    // Get equivalent stresses for elements
    stresses := getEquivalentStresses(els, displacements)

    return displacements, stresses
}