Error prototype

Author: frostedoyster

python/metatensor-learn/metatensor/learn/_backend.py

@@ -9,7 +9,14 @@
 #
 # Any change to this file MUST be also be made to `metatensor/torch/learn.py`.
 
-from metatensor import Labels, LabelsEntry, TensorBlock, TensorMap
+from metatensor import (
+    Labels,
+    LabelsEntry,
+    TensorBlock,
+    TensorMap,
+    equal_metadata,
+    equal_metadata_block,
+)
 
 
 def torch_jit_is_scripting():
@@ -19,6 +26,8 @@ def torch_jit_is_scripting():
 check_isinstance = isinstance
 
 __all__ = [
+    "equal_metadata",
+    "equal_metadata_block",
     "Labels",
     "LabelsEntry",
     "TensorBlock",

python/metatensor-learn/metatensor/learn/nn/errors.py

@@ -2,23 +2,82 @@
 Module containing the :py:class:`AbsoluteError` and :py:class:`SquaredError` classes.
 """
 
-from .._backend import TensorMap
-from .module_map import ModuleMap
+import torch
 
+from .._backend import TensorBlock, TensorMap, equal_metadata, equal_metadata_block
 
-class AbsoluteError(ModuleMap):
 
-    def __init__(self):
-        pass
+def absolute_error(A: TensorMap, B: TensorMap) -> TensorMap:
+    if not equal_metadata(A, B):
+        raise ValueError(
+            "The two maps must have the same metadata in `absolute_error`."
+        )
 
-    def __call__(self, A: TensorMap, B: TensorMap) -> TensorMap:
-        pass
+    keys = []
+    blocks = []
+    for key, block_A in A.items():
+        block_B = B.block(key)
+        keys.append(key)
+        blocks.append(absolute_error_block(block_A, block_B))
 
+    return TensorMap(keys, blocks)
 
-class SquaredError(ModuleMap):
 
-    def __init__(self):
-        pass
+def absolute_error_block(A: TensorBlock, B: TensorBlock) -> TensorBlock:
+    if not equal_metadata_block(A, B):
+        raise ValueError(
+            "The two blocks must have the same metadata in `absolute_error_block`."
+        )
 
-    def __call__(self, A: TensorMap, B: TensorMap) -> TensorMap:
-        pass
+    values = torch.abs(A.values - B.values)
+    block = TensorBlock(
+        values=values,
+        samples=A.samples,
+        components=A.components,
+        properties=A.properties,
+    )
+    for gradient_name, gradient_A in A.gradients():
+        gradient_B = B.gradient(gradient_name)
+        block.add_gradient(
+            gradient_name,
+            absolute_error_block(gradient_A, gradient_B),
+        )
+
+    return block
+
+
+def squared_error(A: TensorMap, B: TensorMap) -> TensorMap:
+    if not equal_metadata(A, B):
+        raise ValueError("The two maps must have the same metadata in `squared_error`.")
+
+    keys = []
+    blocks = []
+    for key, block_A in A.items():
+        block_B = B.block(key)
+        keys.append(key)
+        blocks.append(squared_error_block(block_A, block_B))
+
+    return TensorMap(keys, blocks)
+
+
+def squared_error_block(A: TensorBlock, B: TensorBlock) -> TensorBlock:
+    if not equal_metadata_block(A, B):
+        raise ValueError(
+            "The two blocks must have the same metadata in `squared_error_block`."
+        )
+
+    values = (A.values - B.values) ** 2
+    block = TensorBlock(
+        values=values,
+        samples=A.samples,
+        components=A.components,
+        properties=A.properties,
+    )
+    for gradient_name, gradient_A in A.gradients():
+        gradient_B = B.gradient(gradient_name)
+        block.add_gradient(
+            gradient_name,
+            squared_error_block(gradient_A, gradient_B),
+        )
+
+    return block

python/metatensor-torch/metatensor/torch/learn.py

@@ -4,7 +4,14 @@
 import torch
 
 import metatensor.learn
-from metatensor.torch import Labels, LabelsEntry, TensorBlock, TensorMap
+from metatensor.torch import (
+    Labels,
+    LabelsEntry,
+    TensorBlock,
+    TensorMap,
+    equal_metadata,
+    equal_metadata_block,
+)
 
 
 # ==================================================================================== #
@@ -25,6 +32,8 @@
 module.__dict__["LabelsEntry"] = LabelsEntry
 module.__dict__["TensorBlock"] = TensorBlock
 module.__dict__["TensorMap"] = TensorMap
+module.__dict__["equal_metadata"] = equal_metadata
+module.__dict__["equal_metadata_block"] = equal_metadata_block
 module.__dict__["torch_jit_is_scripting"] = torch.jit.is_scripting