[BUG] Can't create an ElementAssembler on a mixed type mesh imported from meshio

[MechEngrLorenzoFiore]

Description

LinearElasticityElementAssembler.from_mesh(mixed_types_mesh) fails when mixed_types_mesh is a mesh imported from ".inp" format, but only if the number of elements of different types mismatches.

Minimal reproducing example

"""
** single_part_mixed_types_inp_file.inp
*Node
1, 0.0, 0.0, 0
2, 0.0, 1.0, 0
3, 1.0, 0.0, 0
4, 1.0, 1.0, 0
5, 0.5, 1.5, 0
6, 0.0, 2.0, 0
*ELEMENT, TYPE=S4
1, 1, 2, 4, 3
*ELEMENT, TYPE=S3
2, 2, 5, 4
3, 2, 6, 5 
"""

import tensormesh as tm

mixed_types_mesh = tm.Mesh.read("single_part_mixed_types_inp_file.inp", reorder=True)

print(mixed_types_mesh)

mixed_types_mesh.save("single_part_mixed_types_inp_file.vtk")

mixed_types_mesh_form = tm.LinearElasticityElementAssembler.from_mesh(mixed_types_mesh, E=21e9, nu=0.2)

Full traceback / error output

Mesh(
    points: torch.Size([6, 2])
    cells: quad:torch.Size([1, 4]),triangle:torch.Size([2, 3])
    point_data:
    cell_data:
    field_data:
)
Traceback (most recent call last):
  File "/home/user/Simulations/tensormesh_dev/tensormesh_import_from_inp_MWE.py", line 9, in <module>
    mixed_types_mesh_form = tm.LinearElasticityElementAssembler.from_mesh(mixed_types_mesh, E=21e9, nu=0.2)
  File "/home/user/miniconda3/envs/tensormesh/lib/python3.14/site-packages/tensormesh/assemble/element_assembler.py", line 819, in from_mesh
    elem_u, elem_v = torch.stack(elem_u, -1).flatten(), torch.stack(elem_v, -1).flatten() # [num_elements * num_basis * num_basis]
                     ~~~~~~~~~~~^^^^^^^^^^^^
RuntimeError: stack expects each tensor to be equal size, but got [1] at entry 0 and [2] at entry 16

Expected behavior

I would expect the mesh to be imported and the form to be created without any issue even if the mesh has mixed element types.

Environment

torch:       2.12.0+cu130
torch_sla:   0.2.1
tensormesh:  0.1.1
python:      3.14.4
os:          Linux-6.18.33.1-microsoft-standard-WSL2-x86_64-with-glibc2.39
cuda:        n/a

Pre-flight checks

• I searched existing issues and did not find a duplicate.
• I confirmed this is a bug, not a usage question.

[kilickursat]

Thanks for the clear reproducer.

The root cause appears to be in ElementAssembler.from_mesh(), specifically in the edge-precomputation step. The current implementation appends value[:, i] / value[:, j] for each element type and then stacks the entire mixed list:

elem_u.append(value[:, i])
elem_v.append(value[:, j])
...
elem_u, elem_v = torch.stack(elem_u, -1).flatten(), torch.stack(elem_v, -1).flatten()

That assumes every element-type block has the same number of elements. In the reproducer, the imported mesh has one quad and two triangles, so the list contains tensors with lengths 1 and 2, and torch.stack fails.

A robust fix is to build the flattened local (row, col) pairs separately for each element type, then concatenate those flattened arrays only after each block has already been flattened. The per-type arrays can also be reused directly when constructing the projector for each element type.

Suggested patch structure:

elem_u_by_type = {}
elem_v_by_type = {}
elem_u_all = []
elem_v_all = []

for element_type, value in elements.items():
    n_element, n_basis = value.shape

    local_elem_u = (
        value[:, :, None]
        .expand(n_element, n_basis, n_basis)
        .reshape(-1)
    )
    local_elem_v = (
        value[:, None, :]
        .expand(n_element, n_basis, n_basis)
        .reshape(-1)
    )

    elem_u_by_type[element_type] = local_elem_u.cpu().numpy().copy()
    elem_v_by_type[element_type] = local_elem_v.cpu().numpy().copy()

    elem_u_all.append(local_elem_u)
    elem_v_all.append(local_elem_v)

elem_u = torch.cat(elem_u_all).cpu().numpy().copy()
elem_v = torch.cat(elem_v_all).cpu().numpy().copy()

Then, when constructing the projector for each element type, use the corresponding per-type arrays instead of slicing the global arrays with ptr:

for element_type, value in elements.items():
    n_element, n_basis = value.shape

    local_elem_u = elem_u_by_type[element_type]
    local_elem_v = elem_v_by_type[element_type]

    elem_eids = np.array(
        eids_csr[local_elem_u, local_elem_v].copy()
    ).ravel()

    if project == "reduce":
        projector[element_type] = ReduceProjector(
            indices=torch.from_numpy(elem_eids),
            from_shape=(n_element, n_basis, n_basis),
            to_shape=(num_edges,),
        )
    elif project == "sparse":
        projector[element_type] = SparseProjector(
            from_=torch.arange(len(elem_eids), device=points.device),
            to_=torch.from_numpy(elem_eids),
            from_shape=(n_element, n_basis, n_basis),
            to_shape=(num_edges,),
        )

This preserves the expected flattened order for each local element block, while removing the assumption that all element-type blocks have the same number of elements.

A regression test can use the exact mixed .inp reproducer from this issue:

def test_from_mesh_supports_mixed_element_counts_from_inp(tmp_path):
    import tensormesh as tm

    inp = tmp_path / "single_part_mixed_types_inp_file.inp"
    inp.write_text(
        """*Node
1, 0.0, 0.0, 0
2, 0.0, 1.0, 0
3, 1.0, 0.0, 0
4, 1.0, 1.0, 0
5, 0.5, 1.5, 0
6, 0.0, 2.0, 0
*ELEMENT, TYPE=S4
1, 1, 2, 4, 3
*ELEMENT, TYPE=S3
2, 2, 5, 4
3, 2, 6, 5
"""
    )

    mesh = tm.Mesh.read(str(inp), reorder=True)

    assembler = tm.LinearElasticityElementAssembler.from_mesh(
        mesh,
        E=21e9,
        nu=0.2,
    )

    assert set(assembler.element_types) == {"quad", "triangle"}

This should make mixed element-type meshes independent of the number of cells in each block.

[walkerchi]

Fixed and shipped to main via:

• fix: mixed-element-type assembly + hertzian_contact dead code (closes #12, #13) #24 — fix the from_mesh assembly path for mixed element-type meshes (the torch.stack → per-type stack-then-cat change), plus a parametrised regression test covering every mesh × reorder × assembler combination.
• fix(examples): split quads along BL-TR diagonal for plotting #26 — fix the bowtie pattern in the rendered Drucker-Prager / elastic footing examples (quad triangulation must split along the BL-TR diagonal, not BL-TL).

Both landed on main through #27 (commit 81fae3d). Closing — please reopen if you hit any remaining mixed-type case that still fails.