Handle two-dimensional time embeddings

pyhealth/models/unified_embedding.py

@@ -40,6 +40,14 @@ class SinusoidalTimeEmbedding(nn.Module):
     Identical in spirit to the positional encoding in "Attention is All You
     Need" but operating on real-valued timestamps rather than integer positions.
 
+    Examples:
+        >>> embedding = SinusoidalTimeEmbedding(dim=2, max_hours=24.0)
+        >>> output = embedding(torch.tensor([0.0, 6.0]))
+        >>> output.shape
+        torch.Size([2, 2])
+        >>> torch.isfinite(output).all().item()
+        True
+
     Args:
         dim: Output embedding dimension (must be even).
         max_hours: Maximum expected time value in hours.  Values are normalised
@@ -57,7 +65,7 @@ def __init__(self, dim: int, max_hours: float = 720.0):
         self.max_hours = max_hours
         half = dim // 2
         freqs = torch.exp(
-            -math.log(10000.0) * torch.arange(half, dtype=torch.float32) / (half - 1)
+            -math.log(10000.0) * torch.arange(half, dtype=torch.float32) / max(half - 1, 1)
         )
         self.register_buffer("freqs", freqs)  # (dim//2,)
 

tests/core/test_unified_embedding.py

@@ -0,0 +1,29 @@
+import math
+import unittest
+
+import torch
+
+from pyhealth.models import SinusoidalTimeEmbedding
+
+
+class TestSinusoidalTimeEmbedding(unittest.TestCase):
+    def test_dim_two_is_finite(self):
+        embedding = SinusoidalTimeEmbedding(dim=2, max_hours=24.0)
+
+        result = embedding(torch.tensor([0.0, 6.0, 24.0]))
+
+        self.assertEqual(result.shape, (3, 2))
+        self.assertTrue(torch.isfinite(result).all())
+
+    def test_standard_dimension_values_are_unchanged(self):
+        embedding = SinusoidalTimeEmbedding(dim=6, max_hours=24.0)
+        time = torch.tensor([6.0])
+        frequencies = torch.tensor([1.0, 0.01, 0.0001])
+        arguments = time.unsqueeze(-1) / 24.0 * 2 * math.pi * frequencies
+        expected = torch.cat([arguments.sin(), arguments.cos()], dim=-1)
+
+        self.assertTrue(torch.allclose(embedding(time), expected))
+
+
+if __name__ == "__main__":
+    unittest.main()