Add interval argument to LocalCartesian

CHANGELOG.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `interval` argument to `LocalCartesian` transformation ([#7533](https://github.com/pyg-team/pytorch_geometric/pull/7533))
 - Enabled `LinkNeighborLoader` to return number of sampled nodes and edges per hop ([#7516](https://github.com/pyg-team/pytorch_geometric/pull/7516))
 - Added the `HM` personalized fashion recommendation dataset ([#7515](https://github.com/pyg-team/pytorch_geometric/pull/7515))
 - Added the `GraphMixer` model ([#7501](https://github.com/pyg-team/pytorch_geometric/pull/7501))

torch_geometric/transforms/local_cartesian.py

@@ -1,3 +1,5 @@
+from typing import Tuple
+
 import torch
 
 from torch_geometric.data import Data
@@ -10,32 +12,40 @@
 class LocalCartesian(BaseTransform):
     r"""Saves the relative Cartesian coordinates of linked nodes in its edge
     attributes (functional name: :obj:`local_cartesian`). Each coordinate gets
-    *neighborhood-normalized* to the interval :math:`{[0, 1]}^D`.
+    *neighborhood-normalized* to a specified interval
+    (:math:`[0, 1]` by default).
 
     Args:
         norm (bool, optional): If set to :obj:`False`, the output will not be
-            normalized to the interval :math:`{[0, 1]}^D`.
-            (default: :obj:`True`)
+            normalized. (default: :obj:`True`)
         cat (bool, optional): If set to :obj:`False`, all existing edge
             attributes will be replaced. (default: :obj:`True`)
+        interval ((float, float), optional): A tuple specifying the lower and
+            upper bound for normalization. (default: :obj:`(0.0, 1.0)`)
     """
-    def __init__(self, norm: bool = True, cat: bool = True):
+    def __init__(
+            self,
+            norm: bool = True,
+            cat: bool = True,
+            interval: Tuple[float, float] = (0.0, 1.0),
+    ):
         self.norm = norm
         self.cat = cat
+        self.interval = interval
 
     def forward(self, data: Data) -> Data:
         (row, col), pos, pseudo = data.edge_index, data.pos, data.edge_attr
 
         cart = pos[row] - pos[col]
         cart = cart.view(-1, 1) if cart.dim() == 1 else cart
 
-        max_value = scatter(cart.abs(), col, 0, pos.size(0), reduce='max')
-        max_value = max_value.max(dim=-1, keepdim=True)[0]
-
         if self.norm:
-            cart = cart / (2 * max_value[col]) + 0.5
-        else:
-            cart = cart / max_value[col]
+            max_value = scatter(cart.abs(), col, 0, pos.size(0), reduce='max')
+            max_value = max_value.max(dim=-1, keepdim=True)[0]
+
+            length = self.interval[1] - self.interval[0]
+            center = (self.interval[0] + self.interval[1]) / 2
+            cart = length * cart / (2 * max_value[col]) + center
 
         if pseudo is not None and self.cat:
             pseudo = pseudo.view(-1, 1) if pseudo.dim() == 1 else pseudo