# Source code for gemseo.problems.topo_opt.fea_disc

# Copyright 2021 IRT Saint Exupéry, https://www.irt-saintexupery.com
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
#    INITIAL AUTHORS - API and implementation and/or documentation
#        :author: Simone Coniglio
#    OTHER AUTHORS   - MACROSCOPIC CHANGES
"""Finite element analysis (FEA) for 2D topology optimization problems."""

from __future__ import annotations

from typing import TYPE_CHECKING

import scipy
from numpy import arange
from numpy import array
from numpy import atleast_2d
from numpy import kron
from numpy import newaxis
from numpy import ones
from numpy import setdiff1d
from numpy import tile
from numpy import zeros

from gemseo.core.discipline import MDODiscipline

if TYPE_CHECKING:
from collections.abc import Sequence

[docs]
class FininiteElementAnalysis(MDODiscipline):
"""Finite Element Analysis for 2D topology optimization problems.

Take in input the Young Modulus vector E and computes in output the compliance, i.e.
twice the work of external forces.
"""

def __init__(
self,
nu: float = 0.3,
n_x: int = 100,
n_y: int = 100,
f_node: int | Sequence[int] = 101 * 101 - 1,
f_direction: int | Sequence[int] = 1,
f_amplitude: int | Sequence[int] = -1,
fixed_nodes: int | Sequence[int] | None = None,
fixed_dir: int | Sequence[int] | None = None,
name: str | None = None,
) -> None:
"""
Args:
nu: The material Poisson's ratio.
n_x: The number of elements in the x-direction.
n_y: The number of elements in the y-direction.
f_node: The indices of the nodes where the forces are applied.
f_direction: The force direction for each f_node,
either 0 for x or 1 for y.
f_amplitude: The force amplitude for each pair (f_node, f_direction).
fixed_nodes: The indices of the nodes where the structure is clamped.
If None, a default value is used.
fixed_dir: The clamped direction for each node, encode 0 for x and 1 for y.
If None, a default value is used.
name: The name of the discipline.
If None, use the class name.
"""  # noqa: D205 D212
super().__init__(name=name)
if fixed_nodes is None:
fixed_nodes = tile(arange(101), 2)
if fixed_dir is None:
fixed_dir = array([0] * 101 + [1] * 101)
self.N_elements = n_x * n_y
self.N_nodes = (n_x + 1) * (n_y + 1)
self.N_DOFs = 2 * self.N_nodes
self.n_x = n_x
self.n_y = n_y
self.nu = nu
self.E = None
self.KE = None
self.iK = None
self.jK = None
self.edofMat = None
self.freedofs = None
self.f_node = f_node
self.f_direction = f_direction
self.f_amplitude = f_amplitude
self.fixednodes = fixed_nodes
self.fixed_dir = fixed_dir
self.prepare_fea()
self.input_grammar.update_from_names(["E"])
self.output_grammar.update_from_names(["compliance"])
self.default_inputs = {"E": ones(self.N_elements)}

def _run(self) -> None:
em = self.get_inputs_by_name("E")
sk = ((self.KE.flatten()[newaxis]).T * em).flatten(order="F")
k_mat = scipy.sparse.coo_matrix(
(sk, (self.iK, self.jK)), shape=(self.N_DOFs, self.N_DOFs)
).tocsc()
k_mat = k_mat[self.freedofs, :][:, self.freedofs]
u_vec = zeros((self.N_DOFs, 1))
f = zeros((self.N_DOFs, 1))
f[2 * self.f_node + self.f_direction, 0] = self.f_amplitude
u_vec[self.freedofs, 0] = scipy.sparse.linalg.spsolve(
k_mat, f[self.freedofs, 0]
)
# Objective function and sensitivity
ce = ones(self.N_elements)
ce[:] = (
(u_vec[self.edofMat].reshape(self.N_elements, 8) @ self.KE)
* u_vec[self.edofMat].reshape(self.N_elements, 8)
).sum(1)
self.local_data["compliance"] = array([(em * ce).sum()])
self._is_linearized = True
self._init_jacobian()
self.jac["compliance"] = {}
self.jac["compliance"]["E"] = atleast_2d(-ce)

[docs]
def prepare_fea(self) -> None:
"""Prepare the Finite Element Analysis."""
self.KE = self.compute_elementary_stiffeness_matrix()

# FE: Build the index vectors for the for coo matrix format.
edof_mat = zeros((self.N_elements, 8), dtype=int)
for elx in range(self.n_x):
for ely in range(self.n_y):
el = ely + elx * self.n_y
n1 = (self.n_y + 1) * elx + ely
n2 = (self.n_y + 1) * (elx + 1) + ely
edof_mat[el, :] = array([
2 * n1 + 2,
2 * n1 + 3,
2 * n2 + 2,
2 * n2 + 3,
2 * n2,
2 * n2 + 1,
2 * n1,
2 * n1 + 1,
])
self.edofMat = edof_mat
# Construct the index pointers for the coo format
self.iK = kron(edof_mat, ones((8, 1))).flatten()
self.jK = kron(edof_mat, ones((1, 8))).flatten()

# Free DOFs
alldofs = array(range(2 * self.N_nodes))
fixeddofs = 2 * self.fixednodes + self.fixed_dir
self.freedofs = setdiff1d(alldofs, fixeddofs)

[docs]
def compute_elementary_stiffeness_matrix(
self,
) -> None:  # noqa: D205,D212,D415
"""Compute the elementary stiffness matrix of 1x1 quadrilateral elements."""
em = 1.0
k = array([
1 / 2 - self.nu / 6,
1 / 8 + self.nu / 8,
-1 / 4 - self.nu / 12,
-1 / 8 + 3 * self.nu / 8,
-1 / 4 + self.nu / 12,
-1 / 8 - self.nu / 8,
self.nu / 6,
1 / 8 - 3 * self.nu / 8,
])
return (
em
/ (1 - self.nu**2)
* array([
[k[0], k[1], k[2], k[3], k[4], k[5], k[6], k[7]],
[k[1], k[0], k[7], k[6], k[5], k[4], k[3], k[2]],
[k[2], k[7], k[0], k[5], k[6], k[3], k[4], k[1]],
[k[3], k[6], k[5], k[0], k[7], k[2], k[1], k[4]],
[k[4], k[5], k[6], k[7], k[0], k[1], k[2], k[3]],
[k[5], k[4], k[3], k[2], k[1], k[0], k[7], k[6]],
[k[6], k[3], k[4], k[1], k[2], k[7], k[0], k[5]],
[k[7], k[2], k[1], k[4], k[3], k[6], k[5], k[0]],
])
)