Batch D3 and SevenNetD3Model for Torch-Sim interface

[alphalm4]

Batch D3 is implemented

• copying pair_d3_for_ase
• Flatten all input first (escaping pointer-to-pointer-to-pointer...), launch all kernels once, then unflatten
• Unuse AtomicAdd (therefore slower than serial version with small batch size)

SevenNetD3Model in torchsim_d3 module directly use Batch D3.
For reference, I also added SevenNetModel + D3Calculator (serial) version in torchsim module. (will be replaced with above)

For testing, high-precision patch in torch-sim might be required: TorchSim/torch-sim@a858fbd (resolved)

• code review @dambi3613
• Float64 wrapper for double-precision MD computation flow
• speed benchmark
• test scripts (wip)
• enable batch d3 in sevenn inference
• replace SevenNetD3Model with batched D3 version then remove torchsim_d3.py
• overflow issues (at least for maxij term, I feel the issue is resolved but not checked yet)
• Doc

[alphalm4]

Something like this might force Torch-Sim to run in float64 while preserving SevenNet works in float32.

class Float64Wrapper:
    """Wraps a float32 model so torch-sim runs in float64 precision.

    Casts state tensors to float32 before calling the model, then casts
    outputs back to float64. Reports dtype=float64 to torch-sim so all
    integrator arithmetic is done in double precision.
    """

    def __init__(self, model):
        self._model = model
        self._device = model.device
        self._dtype = torch.float64

    @property
    def device(self):
        return self._device

    @property
    def dtype(self):
        return self._dtype

    @property
    def compute_stress(self):
        return self._model.compute_stress

    @property
    def compute_forces(self):
        return self._model.compute_forces

    @property
    def memory_scales_with(self):
        return getattr(self._model, "memory_scales_with", "n_atoms_x_density")

    def __call__(self, state):
        # Cast state to float32 for the model
        state_f32 = state.to(dtype=torch.float32)
        output = self._model(state_f32)
        # Cast outputs back to float64
        return {
            k: v.to(dtype=torch.float64) if isinstance(v, torch.Tensor) else v
            for k, v in output.items()
        }

    def __getattr__(self, name: str):
        return getattr(self._model, name)

[YutackPark]

I dont know the context but why dont we just cast 7net model to float64? For speed?

[alphalm4]

7net model을 64로 올리지않는건 speed 때문이 맞습니다. (굳..이?)

맥락을 추가하자면 해당 항목은 TorchSim state precision이 model precision을 따라가기 때문입니다. SevenNet은 float32 만 받기 때문에 현재 기준으로는 강제로 MD도 float32로 돌려야 합니다.
https://github.com/TorchSim/torch-sim/blob/c456ecec0dec1334a13c026dbcf231fa8309c849/torch_sim/runners.py#L297

아주 정밀하게 확인한 건 아니지만, float64 wrapper 구현 후 ensemble invariant (e.g. npt_nose_hoover_invariant) fluctuation을 확인했을 때 float32보다 줄어드는 것으로 보입니다 (즉 계산 오차 감소가 눈에 보입니다). 반드시 그것뿐만이 아니라도, MD precision을 선택할 방법은 있어야 할 것 같습니다.
물론 TorchSim issue에 올려도 되지만 어렵지않게 구현가능해보여서 넣었습니다.