fem_menos_de_50_lineas.py

import numpy as np
import matplotlib.pyplot as plt


def assem(coords, elems, source):
    """
    Ensambla la matriz de rigidez y el vector de cargas
    
    Parámetros
    ----------
    coords : ndarray, float
        Coordenadas de los nodos.
    elems : ndarray, int
        Conectividad de los elementos.
    source : ndarray, float
        Término fuente evaluado en los nodos de la malla.
    
    Retorna
    -------
    stiff : ndarray, float
        Matriz de rigidez del problema.
    rhs : ndarray, float
        Vector de cargas del problema.
    """
    ncoords = coords.shape[0]
    stiff = np.zeros((ncoords, ncoords))
    rhs = np.zeros((ncoords))
    for el_cont, elem in enumerate(elems):
        stiff_loc, det = local_stiff(coords[elem])
        rhs[elem] += det*np.mean(source[elem])/6
        for row in range(3):
            for col in range(3):
                row_glob, col_glob = elem[row], elem[col]
                stiff[row_glob, col_glob] += stiff_loc[row, col]
    return stiff, rhs


def local_stiff(coords):
    """Calcula la matriz de rigidez local
        
    Parámetros
    ----------
    coords : ndarray, float
        Coordenadas de los nodos del elemento.
    
    Retorna
    -------
    stiff : ndarray, float
        Matriz de rigidez local.
    det : float
        Determinante del jacobian.
    """
    dNdr = np.array([
            [-1, 1, 0],
            [-1, 0, 1]])
    jaco = dNdr @ coords
    det = np.linalg.det(jaco)
    jaco_inv = np.linalg.inv(jaco)
    dNdx = jaco_inv @ dNdr
    stiff = 0.5 * det * (dNdx.T @ dNdx)
    return stiff, det


def pts_restringidos(coords, malla, lineas_rest):
    """
    Identifica los nodos restringidos y libres
    para la malla.
        
    Parámetros
    ----------
    coords : ndarray, float
        Coordenadas de los nodos del elemento.
    malla : Meshio mesh
        Objeto de malla de Meshio.
    lineas_rest : list
        Lista con los números para las líneas
        restringidas.
    
    Retorna
    -------
    pts_rest : list
        Lista de puntos restringidos.
    pts_libres : list
        Lista de puntos libres.
    """
    lineas = [malla.cells[k].data for k in lineas_rest]
    pts_rest = []
    for linea in lineas:
        pts_rest += set(linea.flatten())
    pts_libres = list(set(range(coords.shape[0])) - set(pts_rest))
    return pts_rest, pts_libres


if __name__ == "__main__":
    
    sq3 = np.sqrt(3)
    coords = np.array([
            [sq3, -1],
            [0, 0],
            [2*sq3, 0],
            [0, 2],
            [2*sq3, 2],
            [sq3, 3],
            [sq3, 1]])
    elems = np.array([
            [1, 0, 6],
            [0, 2, 6],
            [2, 4, 6],
            [4, 5, 6],
            [5, 3, 6],
            [3, 1, 6]])

    plt.figure()
    plt.triplot(coords[:, 0], coords[:, 1])
    plt.axis("image")
    plt.show()

    source = np.ones(7)
    stiff, rhs = assem(coords, elems, source)
    free = range(6)
    sol = np.linalg.solve(stiff[np.ix_(free, free)], rhs[free])
    sol_c = np.zeros(coords.shape[0])
    sol_c[free] = sol

    plt.figure()
    plt.tricontourf(coords[:, 0], coords[:, 1], sol_c, 20)
    plt.colorbar()
    plt.axis("image")
    plt.show()