Per-system pbc: batching periodic and non-periodic systems together

[lil-lon]

In the review of #320 there was a short discussion about making pbc a per-system attribute so that systems with different pbc can be mixed in one batch (#320 (comment)).
It was deferred to a separate issue at the time, but I don't think that issue was ever opened, so raising it here.

I think this would be a useful addition. One concern in that discussion was that mixed pbc may complicate the neighbor list calculations (#320 (comment)), but the current neighbor list API already supports per-system pbc (e.g. batch_naive_neighbor_list in nvalchemi-toolkit-ops takes pbc of shape (num_systems, 3): https://github.com/NVIDIA/nvalchemi-toolkit-ops/blob/2b7c3c3adfb1ca84b886eecbf14bc60ff6ba1dc2/nvalchemiops/torch/neighbors/batch_naive.py#L765-L767), and I think other methods support similar APIs, so it can be handled without making the neighbor list logic so complicated.

What do you think? If we agree on a direction I'd be happy to work on a PR.

[CompRhys]

I would be happy to see this PR, I think it's unlikely that the native torch NL we have would support but I also don't believe that they're likely to be used in most cases. Primary question would be how to add the check to the NL that do not support so they raise errors correctly.

[lil-lon]

Thanks! I'd like to start working on a PR. Sharing the rough direction first:

• Make SimState.pbc a per-system (n_systems, 3) tensor. The constructor keeps accepting bool / list[bool] / (3,) and broadcasts, so the main breaking change is the shape of state.pbc for code reading it directly.
• Neighbor lists turned out to need no changes: all backends including the native torch NL already support per-system pbc, and existing tests already cover mixed batches. So I don't think NL-level error checks are needed.
• Core MD and optimization are pbc-agnostic and work with mixed batches as-is. Features that assume periodicity need case-by-case handling (details in the PR): Ewald/PME would zero-fill non-periodic systems per system, matching their current all-non-periodic behavior, and workflows requiring full periodicity (phonons, elastic) get an entry-point check with a clear error.
• In-repo models consume pbc through torchsim_nl and work as-is. MLIP wrappers build their inputs upstream, so mixed-pbc support there is up to each package; I'd provide a small shared helper wrappers can use to require uniform pbc with a clear error.