Do not mix and match metatensor-core and metatensor-torch in gap

Author: Luthaf

src/metatrain/gap/model.py

@@ -6,10 +6,7 @@
 import numpy as np
 import scipy
 import torch
-from metatensor import Labels, TensorBlock, TensorMap
-from metatensor.torch import Labels as TorchLabels
-from metatensor.torch import TensorBlock as TorchTensorBlock
-from metatensor.torch import TensorMap as TorchTensorMap
+from metatensor.torch import Labels, TensorBlock, TensorMap
 from metatomic.torch import (
     AtomisticModel,
     ModelCapabilities,
@@ -121,14 +118,14 @@ def __init__(self, hypers: Dict, dataset_info: DatasetInfo) -> None:
         self._sampler = _FPS(n_to_select=self.hypers["krr"]["num_sparse_points"])
 
         # set it do dummy keys, these are properly set during training
-        self._keys = TorchLabels.empty("_")
+        self._keys = Labels.empty("_")
 
-        dummy_weights = TorchTensorMap(
-            TorchLabels(["_"], torch.tensor([[0]])),
+        dummy_weights = TensorMap(
+            Labels(["_"], torch.tensor([[0]])),
             [mts.block_from_array(torch.empty(1, 1))],
         )
-        dummy_X_pseudo = TorchTensorMap(
-            TorchLabels(["_"], torch.tensor([[0]])),
+        dummy_X_pseudo = TensorMap(
+            Labels(["_"], torch.tensor([[0]])),
             [mts.block_from_array(torch.empty(1, 1))],
         )
         self._subset_of_regressors_torch = TorchSubsetofRegressors(
@@ -138,7 +135,7 @@ def __init__(self, hypers: Dict, dataset_info: DatasetInfo) -> None:
                 "aggregate_names": ["atom", "center_type"],
             },
         )
-        self._species_labels: TorchLabels = TorchLabels.empty("_")
+        self._species_labels: Labels = Labels.empty("_")
 
         # additive models: these are handled by the trainer at training
         # time, and they are added to the output at evaluation time
@@ -186,27 +183,27 @@ def forward(
         self,
         systems: List[System],
         outputs: Dict[str, ModelOutput],
-        selected_atoms: Optional[TorchLabels] = None,
-    ) -> Dict[str, TorchTensorMap]:
+        selected_atoms: Optional[Labels] = None,
+    ) -> Dict[str, TensorMap]:
         soap_features = self._soap_torch_calculator(
             systems, selected_samples=selected_atoms
         )
         # move keys and species labels to device
         self._keys = self._keys.to(systems[0].device)
         self._species_labels = self._species_labels.to(systems[0].device)
 
-        new_blocks: List[TorchTensorBlock] = []
+        new_blocks: List[TensorBlock] = []
         # HACK: to add a block of zeros if there are missing species
         # which were present at training time
         # (with samples "system", "atom" = 0, 0)
         # given the values are all zeros, it does not introduce an error
-        dummyblock: TorchTensorBlock = TorchTensorBlock(
+        dummyblock = TensorBlock(
             values=torch.zeros(
                 (1, len(soap_features[0].properties)),
                 dtype=systems[0].positions.dtype,
                 device=systems[0].device,
             ),
-            samples=TorchLabels(
+            samples=Labels(
                 ["system", "atom"],
                 torch.tensor([[0, 0]], dtype=torch.int, device=systems[0].device),
             ),
@@ -215,7 +212,7 @@ def forward(
         )
         if len(soap_features[0].gradients_list()) > 0:
             for idx, grad in enumerate(soap_features[0].gradients_list()):
-                dummyblock_grad: TorchTensorBlock = TorchTensorBlock(
+                dummyblock_grad = TensorBlock(
                     values=torch.zeros(
                         (
                             1,
@@ -225,7 +222,7 @@ def forward(
                         dtype=systems[0].positions.dtype,
                         device=systems[0].device,
                     ),
-                    samples=TorchLabels(
+                    samples=Labels(
                         ["sample", "system", "atom"],
                         torch.tensor(
                             [[0, 0, 0]], dtype=torch.int, device=systems[0].device
@@ -242,15 +239,15 @@ def forward(
                 new_blocks.append(soap_features.block(key))
             else:
                 new_blocks.append(dummyblock)
-        soap_features = TorchTensorMap(keys=self._species_labels, blocks=new_blocks)
+        soap_features = TensorMap(keys=self._species_labels, blocks=new_blocks)
         soap_features = soap_features.keys_to_samples("center_type")
         # here, we move to properties to use metatensor operations to aggregate
         # later on. Perhaps we could retain the sparsity all the way to the kernels
         # of the soap features with a lot more implementation effort
         soap_features = soap_features.keys_to_properties(
             ["neighbor_1_type", "neighbor_2_type"]
         )
-        soap_features = TorchTensorMap(self._keys, soap_features.blocks())
+        soap_features = TensorMap(self._keys, soap_features.blocks())
         output_key = list(outputs.keys())[0]
         energies = self._subset_of_regressors_torch(soap_features)
         return_dict = {output_key: energies}
@@ -475,9 +472,9 @@ def __init__(
 
     def aggregate_kernel(
         self,
-        kernel: TorchTensorMap,
+        kernel: TensorMap,
         are_pseudo_points: Tuple[bool, bool] = (False, False),
-    ) -> TorchTensorMap:
+    ) -> TensorMap:
         if not are_pseudo_points[0]:
             kernel = mts.sum_over_samples(kernel, self._aggregate_names)
         if not are_pseudo_points[1]:
@@ -488,17 +485,15 @@ def aggregate_kernel(
 
     def forward(
         self,
-        tensor1: TorchTensorMap,
-        tensor2: TorchTensorMap,
+        tensor1: TensorMap,
+        tensor2: TensorMap,
         are_pseudo_points: Tuple[bool, bool] = (False, False),
-    ) -> TorchTensorMap:
+    ) -> TensorMap:
         return self.aggregate_kernel(
             self.compute_kernel(tensor1, tensor2), are_pseudo_points
         )
 
-    def compute_kernel(
-        self, tensor1: TorchTensorMap, tensor2: TorchTensorMap
-    ) -> TorchTensorMap:
+    def compute_kernel(self, tensor1: TensorMap, tensor2: TensorMap) -> TensorMap:
         raise NotImplementedError("compute_kernel needs to be implemented.")
 
 
@@ -512,7 +507,7 @@ def __init__(
         super().__init__(aggregate_names, structurewise_aggregate)
         self._degree = degree
 
-    def compute_kernel(self, tensor1: TorchTensorMap, tensor2: TorchTensorMap):
+    def compute_kernel(self, tensor1: TensorMap, tensor2: TensorMap):
         return mts.pow(mts.dot(tensor1, tensor2), self._degree)
 
 
@@ -546,10 +541,6 @@ def fit(self, X: TensorMap):  # -> GreedySelector:
         :param X:
             Training vectors.
         """
-        if isinstance(X, torch.ScriptObject):
-            X = torch_tensor_map_to_core(X)
-            assert isinstance(X[0].values, np.ndarray)
-
         if len(X.component_names) != 0:
             raise ValueError("Only blocks with no components are supported.")
 
@@ -578,7 +569,9 @@ def fit(self, X: TensorMap):  # -> GreedySelector:
 
             blocks.append(
                 TensorBlock(
-                    values=np.zeros([len(samples), len(properties)], dtype=np.int32),
+                    values=torch.zeros(
+                        [len(samples), len(properties)], dtype=torch.int32
+                    ),
                     samples=samples,
                     components=[],
                     properties=properties,
@@ -596,12 +589,6 @@ def transform(self, X: TensorMap) -> TensorMap:
         :returns:
             The selected subset of the input.
         """
-        if isinstance(X, torch.ScriptObject):
-            use_mts_torch = True
-            X = torch_tensor_map_to_core(X)
-        else:
-            use_mts_torch = False
-
         blocks = []
         for key, block in X.items():
             block_support = self.support.block(key)
@@ -614,10 +601,7 @@ def transform(self, X: TensorMap) -> TensorMap:
                 new_block = mts.slice_block(block, "samples", block_support.samples)
             blocks.append(new_block)
 
-        X_reduced = TensorMap(X.keys, blocks)
-        if use_mts_torch:
-            X_reduced = core_tensor_map_to_torch(X_reduced)
-        return X_reduced
+        return TensorMap(X.keys, blocks)
 
     def fit_transform(self, X: TensorMap) -> TensorMap:
         """Fit to data, then transform it.
@@ -628,112 +612,6 @@ def fit_transform(self, X: TensorMap) -> TensorMap:
         return self.fit(X).transform(X)
 
 
-def torch_tensor_map_to_core(torch_tensor: TorchTensorMap):
-    torch_blocks = []
-    for _, torch_block in torch_tensor.items():
-        torch_blocks.append(torch_tensor_block_to_core(torch_block))
-    torch_keys = torch_labels_to_core(torch_tensor.keys)
-    return TensorMap(torch_keys, torch_blocks)
-
-
-def torch_tensor_block_to_core(torch_block: TorchTensorBlock):
-    """Transforms a tensor block from metatensor-torch to metatensor-torch
-    :param torch_block:
-        tensor block from metatensor-torch
-    :returns torch_block:
-        tensor block from metatensor-torch
-    """
-    block = TensorBlock(
-        values=torch_block.values.detach().cpu().numpy(),
-        samples=torch_labels_to_core(torch_block.samples),
-        components=[
-            torch_labels_to_core(component) for component in torch_block.components
-        ],
-        properties=torch_labels_to_core(torch_block.properties),
-    )
-    for parameter, gradient in torch_block.gradients():
-        block.add_gradient(
-            parameter=parameter,
-            gradient=TensorBlock(
-                values=gradient.values.detach().cpu().numpy(),
-                samples=torch_labels_to_core(gradient.samples),
-                components=[
-                    torch_labels_to_core(component) for component in gradient.components
-                ],
-                properties=torch_labels_to_core(gradient.properties),
-            ),
-        )
-    return block
-
-
-def torch_labels_to_core(torch_labels: TorchLabels):
-    """Transforms labels from metatensor-torch to metatensor-torch
-    :param torch_block:
-        tensor block from metatensor-torch
-    :returns torch_block:
-        labels from metatensor-torch
-    """
-    return Labels(torch_labels.names, torch_labels.values.detach().cpu().numpy())
-
-
-###
-
-
-def core_tensor_map_to_torch(core_tensor: TensorMap):
-    """Transforms a tensor map from metatensor-core to metatensor-torch
-    :param core_tensor:
-        tensor map from metatensor-core
-    :returns torch_tensor:
-        tensor map from metatensor-torch
-    """
-
-    torch_blocks = []
-    for _, core_block in core_tensor.items():
-        torch_blocks.append(core_tensor_block_to_torch(core_block))
-    torch_keys = core_labels_to_torch(core_tensor.keys)
-    return TorchTensorMap(torch_keys, torch_blocks)
-
-
-def core_tensor_block_to_torch(core_block: TensorBlock):
-    """Transforms a tensor block from metatensor-core to metatensor-torch
-    :param core_block:
-        tensor block from metatensor-core
-    :returns torch_block:
-        tensor block from metatensor-torch
-    """
-    block = TorchTensorBlock(
-        values=torch.tensor(core_block.values),
-        samples=core_labels_to_torch(core_block.samples),
-        components=[
-            core_labels_to_torch(component) for component in core_block.components
-        ],
-        properties=core_labels_to_torch(core_block.properties),
-    )
-    for parameter, gradient in core_block.gradients():
-        block.add_gradient(
-            parameter=parameter,
-            gradient=TorchTensorBlock(
-                values=torch.tensor(gradient.values),
-                samples=core_labels_to_torch(gradient.samples),
-                components=[
-                    core_labels_to_torch(component) for component in gradient.components
-                ],
-                properties=core_labels_to_torch(gradient.properties),
-            ),
-        )
-    return block
-
-
-def core_labels_to_torch(core_labels: Labels):
-    """Transforms labels from metatensor-core to metatensor-torch
-    :param core_block:
-        tensor block from metatensor-core
-    :returns torch_block:
-        labels from metatensor-torch
-    """
-    return TorchLabels(core_labels.names, torch.tensor(core_labels.values))
-
-
 class SubsetOfRegressors:
     def __init__(
         self,
@@ -809,10 +687,6 @@ def fit(
             if not isinstance(alpha_forces, float):
                 raise ValueError("alpha must either be a float")
 
-        X = X.to(arrays="numpy")
-        X_pseudo = X_pseudo.to(arrays="numpy")
-        y = y.to(arrays="numpy")
-
         if self._kernel is None:
             # _set_kernel only returns None if kernel type is precomputed
             k_nm = X
@@ -831,11 +705,11 @@ def fit(
             structures = torch.unique(k_nm_block.samples["system"])
             n_atoms_per_structure = []
             for structure in structures:
-                n_atoms = np.sum(X_block.samples["system"] == structure)
+                n_atoms = torch.sum(X_block.samples["system"] == structure)
                 n_atoms_per_structure.append(float(n_atoms))
 
-            n_atoms_per_structure = np.array(n_atoms_per_structure)
-            normalization = np.sqrt(n_atoms_per_structure)
+            n_atoms_per_structure = torch.tensor(n_atoms_per_structure)
+            normalization = torch.sqrt(n_atoms_per_structure)
 
             if not (np.allclose(alpha_energy, 0.0)):
                 normalization /= alpha_energy
@@ -871,7 +745,7 @@ def fit(
             self._solver.fit(k_nm_reg, y_reg)
 
             weight_block = TensorBlock(
-                values=self._solver.weights.T,
+                values=torch.as_tensor(self._solver.weights.T),
                 samples=y_block.properties,
                 components=k_nm_block.components,
                 properties=k_nm_block.properties,
@@ -901,17 +775,17 @@ def predict(self, T: TensorMap) -> TensorMap:
 
     def export_torch_script_model(self):
         return TorchSubsetofRegressors(
-            core_tensor_map_to_torch(self._weights),
-            core_tensor_map_to_torch(self._X_pseudo),
+            self._weights,
+            self._X_pseudo,
             self._kernel_kwargs,
         )
 
 
 class TorchSubsetofRegressors(torch.nn.Module):
     def __init__(
         self,
-        weights: TorchTensorMap,
-        X_pseudo: TorchTensorMap,
+        weights: TensorMap,
+        X_pseudo: TensorMap,
         kernel_kwargs: Optional[dict] = None,
     ):
         super().__init__()
@@ -923,7 +797,7 @@ def __init__(
         # Set the kernel
         self._kernel = TorchAggregatePolynomial(**kernel_kwargs)
 
-    def forward(self, T: TorchTensorMap) -> TorchTensorMap:
+    def forward(self, T: TensorMap) -> TensorMap:
         """
         :param T:
             features