Implement SimState.wrap_positions using pbc_wrap_batched

[Praagnya]

Summary

• Implemented SimState.wrap_positions property using ts.transforms.pbc_wrap_batched to properly wrap atomic positions according to periodic boundary conditions.

Checklist

Before a pull request can be merged, the following items must be checked:

• Doc strings have been added in the Google docstring format.
• Run ruff on your code.
• Tests have been added for any new functionality or bug fixes.

Resolves #343

[Praagnya]

@thomasloux Could you please review this? I created a new branch without the test since my previous commits got messy after other changes in my fork. So I opened a clean PR.

[orionarcher]

Could you link the other PR for context?

[Praagnya]

Previous PR: #345

[thomasloux]

now self.pbc is a Boolean (3,) shaped Tensor. So this expression won't work here. I don't think this piece of code should need a test, a static typing check would have catched that. replace with if not self.pbc.any()

[thomasloux]

For reference of use of static check #339