Fix unit_cell_fire RuntimeError from bad dtype (#27)

* fix RuntimeError: scatter(): Expected self.dtype to be equal to src.dtype

torchsim/optimizers.py", line 768, in fire_step
 atomic_power_per_batch.scatter_add_(
RuntimeError: scatter(): Expected self.dtype to be equal to src.dtype

* fix MaceModel init by setting self.pbc_template ahead of first potential use in setup_from_batch

Author: janosh

torchsim/models/mace.py

@@ -94,7 +94,7 @@ def __init__(
         )
 
         # setup system boundary conditions
-        pbc = [True] * 3 if periodic else [False] * 3
+        pbc = [periodic] * 3
         self.pbc = torch.tensor([pbc], device=self.device)
 
         if atomic_numbers is not None:
@@ -331,6 +331,11 @@ def __init__(
         # Store flag to track if atomic numbers were provided at init
         self.atomic_numbers_in_init = atomic_numbers is not None
 
+        # Set PBC
+        pbc = [periodic] * 3
+        self.pbc_template = torch.tensor([pbc], device=self._device)
+        self.pbc = None  # Will be set in forward
+
         # Set up batch information if atomic numbers are provided
         if atomic_numbers is not None:
             if batch is None:
@@ -341,11 +346,6 @@ def __init__(
 
             self.setup_from_batch(atomic_numbers, batch)
 
-        # Set PBC
-        pbc = [True] * 3 if periodic else [False] * 3
-        self.pbc_template = torch.tensor([pbc], device=self._device)
-        self.pbc = None  # Will be set in forward
-
     def setup_from_batch(self, atomic_numbers: torch.Tensor, batch: torch.Tensor) -> None:
         """Set up internal state from atomic numbers and batch indices.
 

torchsim/optimizers.py

@@ -764,7 +764,9 @@ def fire_step(  # noqa: PLR0915
 
         # Calculate power (F·V) for atoms
         atomic_power = (state.forces * state.velocity).sum(dim=1)  # [n_atoms]
-        atomic_power_per_batch = torch.zeros(n_batches, device=device, dtype=dtype)
+        atomic_power_per_batch = torch.zeros(
+            n_batches, device=device, dtype=atomic_power.dtype
+        )
         atomic_power_per_batch.scatter_add_(
             dim=0, index=state.batch, src=atomic_power
         )  # [n_batches]
@@ -774,7 +776,9 @@ def fire_step(  # noqa: PLR0915
             dim=1
         )  # [n_batches*3]
         cell_batch = torch.arange(n_batches, device=device).repeat_interleave(3)
-        cell_power_per_batch = torch.zeros(n_batches, device=device, dtype=dtype)
+        cell_power_per_batch = torch.zeros(
+            n_batches, device=device, dtype=cell_power.dtype
+        )
         cell_power_per_batch.scatter_add_(
             dim=0, index=cell_batch, src=cell_power
         )  # [n_batches]