Skip to content

[Code scan] Fix AtomicData registration and batched conversion crashes #38

Description

@njzjz

This issue comes from a Codex global repository scan.

Problem

Several AtomicData code paths mix PyTorch assumptions into NumPy objects or mishandle optional batched data:

register_fields(env_fields=...) and register_fields(onsitenv_fields=...) always fail because allfields includes those fields but the RHS uniqueness count omits them:

allfields = node_fields.union(edge_fields, graph_fields, env_fields, onsitenv_fields)
assert len(allfields) == len(node_fields) + len(edge_fields) + len(graph_fields)

Batched validation uses ndarray.view(-1, 3, 3), which is not reshape for NumPy arrays:

cell = self.cell.view(-1, 3, 3)

to_ase() enters the cell is None branch and then dereferences cell.shape:

if batch is not None:
n_batches = batch.max() + 1
if cell is None:
if cell.shape[0] == 1:
cell = np.repeat(cell, n_batches, axis=0)

Unbatched to_ase(extra_fields=...) sets masks to slice(None) and later calls .sum() on those slices:

self[key][mask].reshape(mask.sum(), -1)
)
elif key in _EDGE_FIELDS:
mol.info[key] = (
self[key][edge_mask].reshape(edge_mask.sum(), -1)

Suggested fix

Include env/onsitenv fields in the uniqueness count and duplicate checks, replace NumPy view(-1, 3, 3) with reshape(-1, 3, 3), handle cell is None without dereferencing it, and branch unbatched extra-field masking away from .sum() on slice objects.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    Status
    Todo

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions