Skip to content

fem

FEM

Bases: MessagePassing

Message-passing layer performing one Jacobi-style equilibrium update of a pin-jointed truss.

Each call aggregates the axial member forces implied by the current displacements into a per-node residual force, then applies a diagonal (per-node) stiffness correction to obtain updated displacements.

Source code in src/torch_structure/message_passing/fem.py 
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
class FEM(MessagePassing):
    """Message-passing layer performing one Jacobi-style equilibrium update of a pin-jointed truss.

    Each call aggregates the axial member forces implied by the current
    displacements into a per-node residual force, then applies a diagonal
    (per-node) stiffness correction to obtain updated displacements.
    """

    def __init__(self):
        super().__init__(aggr="add")

    def forward(
        self,
        x,
        load,
        u,
        support,
        effective_stiffness,
        edge_index,
        length,
        A,
        E,
        direction,
        transformation_matrix,
    ):
        """Perform one equilibrium-update iteration.

        Args:
            x (torch.Tensor [N, 3]): node coordinates, passed through for
                propagate's argument resolution but not used by ``message``.
            load (torch.Tensor [N, 3]): externally applied nodal loads.
            u (torch.Tensor [N, 3]): current nodal displacements.
            support (torch.Tensor [N], bool): mask of supported (fixed) nodes,
                whose displacement is reset to zero after the update.
            effective_stiffness (torch.Tensor [N]): per-node diagonal
                stiffness used to convert residual force into a displacement
                correction.
            edge_index (torch.Tensor [2, E]): edge connectivity.
            length (torch.Tensor [E]): member lengths.
            A (torch.Tensor [E]): member cross-sectional areas.
            E (torch.Tensor [E]): member Young's moduli.
            direction (torch.Tensor [E, 3]): unit vector along each member.
            transformation_matrix (torch.Tensor [E, 2, 6]): local
                displacement transformation matrix for each member.

        Returns:
            tuple[torch.Tensor, torch.Tensor]: the updated nodal
            displacements ``u_new`` and the per-node residual force before
            the update.
        """
        residual_force = load + self.propagate(
            edge_index,
            x=x,
            u=u,
            length=length,
            A=A,
            E=E,
            direction=direction,
            transformation_matrix=transformation_matrix,
        )
        effective_stiffness = effective_stiffness.view(-1, 1) + 1e-9
        u_new = u + residual_force / effective_stiffness
        u_new[support] = 0
        return u_new, residual_force

    def message(self, u_i, u_j, length, A, E, direction, transformation_matrix):
        """Compute the axial member force vector contributed to the target (center) node.

        Transforms the target/source displacements into local axial
        coordinates, derives elongation, strain and stress from the member's
        length and Young's modulus, and returns the resulting force vector
        along the member direction.
        """
        length = length.view(-1)
        A = A.view(-1)
        E = E.view(-1)
        local_u = torch.matmul(
            transformation_matrix, torch.cat([u_i, u_j], dim=1).unsqueeze(2)
        ).squeeze(2)
        elongation = local_u[:, 0] - local_u[:, 1]
        strain = elongation / length
        stress = E * strain
        force = stress * A
        force_vector = -1 * direction * force.view(-1, 1)
        return force_vector

forward(x, load, u, support, effective_stiffness, edge_index, length, A, E, direction, transformation_matrix)

Perform one equilibrium-update iteration.

Parameters:

Name Type Description Default
x Tensor[N, 3]

node coordinates, passed through for propagate's argument resolution but not used by message.

 required
load Tensor[N, 3]

externally applied nodal loads.

 required
u Tensor[N, 3]

current nodal displacements.

 required
support (Tensor[N], bool)

mask of supported (fixed) nodes, whose displacement is reset to zero after the update.

 required
effective_stiffness Tensor[N]

per-node diagonal stiffness used to convert residual force into a displacement correction.

 required
edge_index Tensor[2, E]

edge connectivity.

 required
length Tensor[E]

member lengths.

 required
A Tensor[E]

member cross-sectional areas.

 required
E Tensor[E]

member Young's moduli.

 required
direction Tensor[E, 3]

unit vector along each member.

 required
transformation_matrix Tensor[E, 2, 6]

local displacement transformation matrix for each member.

 required

Returns:

Type Description

tuple[torch.Tensor, torch.Tensor]: the updated nodal

displacements u_new and the per-node residual force before

the update.
Source code in src/torch_structure/message_passing/fem.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def forward(
    self,
    x,
    load,
    u,
    support,
    effective_stiffness,
    edge_index,
    length,
    A,
    E,
    direction,
    transformation_matrix,
):
    """Perform one equilibrium-update iteration.

    Args:
        x (torch.Tensor [N, 3]): node coordinates, passed through for
            propagate's argument resolution but not used by ``message``.
        load (torch.Tensor [N, 3]): externally applied nodal loads.
        u (torch.Tensor [N, 3]): current nodal displacements.
        support (torch.Tensor [N], bool): mask of supported (fixed) nodes,
            whose displacement is reset to zero after the update.
        effective_stiffness (torch.Tensor [N]): per-node diagonal
            stiffness used to convert residual force into a displacement
            correction.
        edge_index (torch.Tensor [2, E]): edge connectivity.
        length (torch.Tensor [E]): member lengths.
        A (torch.Tensor [E]): member cross-sectional areas.
        E (torch.Tensor [E]): member Young's moduli.
        direction (torch.Tensor [E, 3]): unit vector along each member.
        transformation_matrix (torch.Tensor [E, 2, 6]): local
            displacement transformation matrix for each member.

    Returns:
        tuple[torch.Tensor, torch.Tensor]: the updated nodal
        displacements ``u_new`` and the per-node residual force before
        the update.
    """
    residual_force = load + self.propagate(
        edge_index,
        x=x,
        u=u,
        length=length,
        A=A,
        E=E,
        direction=direction,
        transformation_matrix=transformation_matrix,
    )
    effective_stiffness = effective_stiffness.view(-1, 1) + 1e-9
    u_new = u + residual_force / effective_stiffness
    u_new[support] = 0
    return u_new, residual_force

message(u_i, u_j, length, A, E, direction, transformation_matrix)

Compute the axial member force vector contributed to the target (center) node.

Transforms the target/source displacements into local axial coordinates, derives elongation, strain and stress from the member's length and Young's modulus, and returns the resulting force vector along the member direction.

Source code in src/torch_structure/message_passing/fem.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def message(self, u_i, u_j, length, A, E, direction, transformation_matrix):
    """Compute the axial member force vector contributed to the target (center) node.

    Transforms the target/source displacements into local axial
    coordinates, derives elongation, strain and stress from the member's
    length and Young's modulus, and returns the resulting force vector
    along the member direction.
    """
    length = length.view(-1)
    A = A.view(-1)
    E = E.view(-1)
    local_u = torch.matmul(
        transformation_matrix, torch.cat([u_i, u_j], dim=1).unsqueeze(2)
    ).squeeze(2)
    elongation = local_u[:, 0] - local_u[:, 1]
    strain = elongation / length
    stress = E * strain
    force = stress * A
    force_vector = -1 * direction * force.view(-1, 1)
    return force_vector

StiffnessAggregator

Bases: MessagePassing

Sums each edge's stiffness contribution onto its target node.

Source code in src/torch_structure/message_passing/fem.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
class StiffnessAggregator(MessagePassing):
    """Sums each edge's stiffness contribution onto its target node."""

    def __init__(self):
        super().__init__(aggr="add")

    def forward(self, edge_index, stiffness):
        """Aggregate per-edge stiffness values into a per-node effective stiffness.

        Args:
            edge_index (torch.Tensor [2, E]): edge connectivity.
            stiffness (torch.Tensor [E]): per-edge stiffness values.

        Returns:
            torch.Tensor [N]: effective stiffness at each node.
        """
        return self.propagate(edge_index, stiffness=stiffness)

    def message(self, stiffness):
        """Pass each edge's stiffness value through unchanged."""
        return stiffness

forward(edge_index, stiffness)

Aggregate per-edge stiffness values into a per-node effective stiffness.

Parameters:

Name Type Description Default
edge_index Tensor[2, E]

edge connectivity.

 required
stiffness Tensor[E]

per-edge stiffness values.

 required

Returns:

Type Description

torch.Tensor [N]: effective stiffness at each node.
Source code in src/torch_structure/message_passing/fem.py 
 98
 99
100
101
102
103
104
105
106
107
108
def forward(self, edge_index, stiffness):
    """Aggregate per-edge stiffness values into a per-node effective stiffness.

    Args:
        edge_index (torch.Tensor [2, E]): edge connectivity.
        stiffness (torch.Tensor [E]): per-edge stiffness values.

    Returns:
        torch.Tensor [N]: effective stiffness at each node.
    """
    return self.propagate(edge_index, stiffness=stiffness)

message(stiffness)

Pass each edge's stiffness value through unchanged.

Source code in src/torch_structure/message_passing/fem.py 
110
111
112
def message(self, stiffness):
    """Pass each edge's stiffness value through unchanged."""
    return stiffness