Add method to add new "atlases"

src/neuromaps_prime/graph/core.py

@@ -115,10 +115,37 @@ def add_transform(
             weight=transform.weight,
         )
 
+    def add_atlas(self, atlas: SurfaceAtlas | VolumeAtlas) -> None:
+        """Register an atlas to a graph node and a cache entry.
+
+        Args:
+            atlas: SurfaceAtlas or VolumeAtlas to add to existing node.
+
+        Raises:
+            TypeError: If atlas is not SurfaceAtlas or VolumeAtlas.
+            ValueError: If atlas space is not present in the graph.
+        """
+        if not isinstance(atlas, (SurfaceAtlas, VolumeAtlas)):
+            raise TypeError(f"Unsupported atlas type: {type(atlas)}")
+
+        node_name = atlas.space
+        if node_name not in self.nodes:
+            raise ValueError(
+                f"Node '{node_name}' not found. Available nodes: {sorted(self.nodes)}"
+            )
+
+        node_data = self.nodes[node_name]["data"]
+        match atlas:
+            case SurfaceAtlas():
+                node_data.surfaces.append(atlas)
+                self._cache.add_surface_atlas(atlas)
+            case VolumeAtlas():
+                node_data.volumes.append(atlas)
+                self._cache.add_volume_atlas(atlas)
+
     # ------------------------------------------------------------------ #
     # Validation                                                           #
     # ------------------------------------------------------------------ #
-
     def validate_spaces(self, source: str, target: str) -> None:
         """Assert that both source and target exist as nodes in the graph.
 

tests/graph/test_graph.py

@@ -140,6 +140,66 @@ def test_add_invalid_transform(self, graph: NeuromapsGraph) -> None:
         with pytest.raises(TypeError, match="Unsupported transform type"):
             graph.add_transform("invalid", key="key")  # type: ignore[arg-type]
 
+    def test_add_missing_atlas_space(self, graph: NeuromapsGraph) -> None:
+        """Test adding atlas with an invalid space."""
+        atlas = models.SurfaceAtlas(
+            name="Test",
+            description="Test",
+            file_path=Path("."),
+            space="invalid",
+            density="10k",
+            hemisphere="left",
+            resource_type="test",
+        )
+        with pytest.raises(ValueError, match="not found"):
+            graph.add_atlas(atlas)
+
+    def test_add_invalid_atlas_type(self, graph: NeuromapsGraph) -> None:
+        """Test adding invalid atlas raises error."""
+        with pytest.raises(TypeError, match="Unsupported atlas type"):
+            graph.add_atlas("invalid")  # type: ignore [arg-type]
+
+    def test_add_surface_atlas(self, graph: NeuromapsGraph) -> None:
+        """Test adding surface atlas."""
+        atlas = models.SurfaceAtlas(
+            name="Test",
+            description="Test",
+            file_path=Path("."),
+            space="Yerkes19",
+            density="10k",
+            hemisphere="left",
+            resource_type="test",
+        )
+        graph.add_atlas(atlas)
+        assert atlas in graph.nodes["Yerkes19"]["data"].surfaces
+        assert atlas not in graph.nodes["Yerkes19"]["data"].volumes
+        assert (
+            graph._cache.get_surface_atlas(
+                atlas.space, atlas.density, atlas.hemisphere, atlas.resource_type
+            )
+            == atlas
+        )
+
+    def test_add_volume_atlas(self, graph: NeuromapsGraph) -> None:
+        """Test adding volume atlas."""
+        atlas = models.VolumeAtlas(
+            name="Test",
+            description="Test",
+            file_path=Path("."),
+            space="Yerkes19",
+            resolution="1mm",
+            resource_type="test",
+        )
+        graph.add_atlas(atlas)
+        assert atlas not in graph.nodes["Yerkes19"]["data"].surfaces
+        assert atlas in graph.nodes["Yerkes19"]["data"].volumes
+        assert (
+            graph._cache.get_volume_atlas(
+                atlas.space, atlas.resolution, atlas.resource_type
+            )
+            == atlas
+        )
+
     def test_fetch_volume_atlas(self, graph: NeuromapsGraph) -> None:
         """Test fetching volume atlas."""
         atlas = graph.fetch_volume_atlas(