Add merge_networks (closes #221)

alchemiscale/interface/api.py

@@ -31,7 +31,7 @@
 from ..settings import get_base_api_settings
 from ..storage.statestore import Neo4jStore
 from ..storage.objectstore import S3ObjectStore
-from ..storage.models import TaskStatusEnum, StrategyState
+from ..storage.models import NetworkStateEnum, TaskStatusEnum, StrategyState
 from ..models import Scope, ScopedKey
 from ..security.models import TokenData, CredentialedUserIdentity
 
@@ -138,6 +138,55 @@ async def create_network(
     return an_sk
 
 
+@router.post("/networks/merge", response_model=ScopedKey)
+async def merge_networks(
+    *,
+    networks: list[str] = Body(embed=True),
+    name: str = Body(embed=True),
+    scope: dict = Body(embed=True),
+    state: str = Body(embed=True, default=NetworkStateEnum.active.value),
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+    token: TokenData = Depends(get_token_data_depends),
+):
+    # validate the destination scope first
+    try:
+        target_scope = Scope(**scope)
+    except (TypeError, ValidationError) as e:
+        raise HTTPException(
+            status_code=http_status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=str(e),
+        )
+    validate_scopes(target_scope, token)
+
+    # validate each source network's scope is accessible to the token
+    network_sks = []
+    for network in networks:
+        try:
+            network_sk = ScopedKey.from_str(network)
+        except ValueError as e:
+            raise HTTPException(
+                status_code=http_status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=e.args[0],
+            )
+        validate_scopes(network_sk.scope, token)
+        network_sks.append(network_sk)
+
+    try:
+        an_sk = n4js.merge_networks(
+            network_scoped_keys=network_sks,
+            name=name,
+            scope=target_scope,
+            state=state,
+        )
+    except ValueError as e:
+        raise HTTPException(
+            status_code=http_status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=e.args[0],
+        )
+
+    return an_sk
+
+
 @router.post("/bulk/networks/state/set")
 def set_networks_state(
     *,

alchemiscale/interface/client.py

@@ -187,6 +187,103 @@ def post():
 
         return ScopedKey.from_dict(scoped_key)
 
+    def merge_networks(
+        self,
+        networks: list[ScopedKey],
+        name: str,
+        scope: Scope,
+        state: NetworkStateEnum | str = NetworkStateEnum.active,
+        visualize: bool = True,
+    ) -> ScopedKey:
+        """Merge multiple existing AlchemicalNetworks into a new AlchemicalNetwork.
+
+        The resulting AlchemicalNetwork contains the union of all
+        Transformations and NonTransformations from the source networks.
+        Existing Tasks for those transformations that are in ``complete`` or
+        ``error`` state are cloned into the new network's scope along with
+        their associated ProtocolDAGResultRefs, so previously-computed results
+        do not need to be re-run.
+
+        Cloned Tasks are wired to their Transformations via ``PERFORMS`` and
+        are reachable through standard network traversals
+        (``get_network_tasks``, ``get_network_results``, etc.). They are
+        intentionally **not** actioned to the new network's TaskHub; to
+        retry errored Tasks on the merged network, call
+        :meth:`action_tasks` with the merged network's ScopedKey after the
+        merge completes.
+
+        Parameters
+        ----------
+        networks
+            The ScopedKeys of the AlchemicalNetworks to merge. The source
+            networks may live in different Scopes; the caller must have access
+            to each.
+        name
+            The name of the new AlchemicalNetwork.
+        scope
+            The Scope in which to create the new AlchemicalNetwork.
+            This must be a *specific* Scope; it must not contain wildcards.
+        state
+            The starting state of the new AlchemicalNetwork in the database.
+            See :meth:`AlchemiscaleClient.set_network_state` for valid states.
+            Defaults to ``"active"``.
+        visualize
+            If ``True``, show submission progress indicator.
+
+        Returns
+        -------
+        ScopedKey
+            The ScopedKey of the new, merged AlchemicalNetwork.
+        """
+        if not scope.specific():
+            raise ValueError(
+                f"`scope` '{scope}' contains wildcards ('*'); `scope` must be *specific*"
+            )
+
+        if not networks:
+            raise ValueError("`networks` must contain at least one ScopedKey")
+
+        network_sks = [
+            sk if isinstance(sk, ScopedKey) else ScopedKey.from_str(sk)
+            for sk in networks
+        ]
+
+        for network_sk in network_sks:
+            if network_sk.qualname not in ("AlchemicalNetwork",):
+                raise ValueError(
+                    f"ScopedKey '{network_sk}' does not refer to an AlchemicalNetwork"
+                )
+
+        state = NetworkStateEnum(state)
+
+        data = dict(
+            networks=[str(sk) for sk in network_sks],
+            name=name,
+            scope=scope.to_dict(),
+            state=state.value,
+        )
+
+        def post():
+            return self._post_resource("/networks/merge", data)
+
+        if visualize:
+            from rich.progress import Progress
+
+            with Progress(*self._rich_waiting_columns(), transient=False) as progress:
+                task = progress.add_task(
+                    f"Merging [bold]{len(network_sks)}[/bold] networks into "
+                    f"[bold]'{name}'[/bold] in scope [bold]'{scope}'[/bold]...",
+                    total=None,
+                )
+
+                scoped_key = post()
+                progress.start_task(task)
+                progress.update(task, total=1, completed=1)
+        else:
+            scoped_key = post()
+
+        return ScopedKey.from_dict(scoped_key)
+
     def set_network_state(
         self, network: ScopedKey, state: NetworkStateEnum | str
     ) -> ScopedKey | None: