Add copy_network, merge_networks, merge_scopes (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,101 @@ 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("/networks/{network_scoped_key}/copy", response_model=ScopedKey)
+async def copy_network(
+    network_scoped_key,
+    *,
+    scope: dict = Body(embed=True),
+    name: str | None = Body(embed=True, default=None),
+    state: str = Body(embed=True, default=NetworkStateEnum.active.value),
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+    token: TokenData = Depends(get_token_data_depends),
+):
+    # validate the source ScopedKey + its scope
+    try:
+        source_sk = ScopedKey.from_str(network_scoped_key)
+    except ValueError as e:
+        raise HTTPException(
+            status_code=http_status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=e.args[0],
+        )
+    validate_scopes(source_sk.scope, token)
+
+    # validate the destination scope
+    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)
+
+    try:
+        an_sk = n4js.copy_network(
+            network_scoped_key=source_sk,
+            scope=target_scope,
+            name=name,
+            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,253 @@ 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.
+        Only Tasks in ``complete`` state on those Transformations are
+        cloned into the new network's scope along with their
+        ProtocolDAGResultRefs, so previously-computed results do not need
+        to be re-run.
+
+        Any ``Strategy`` (:meth:`set_network_strategy`) and
+        ``TaskRestartPattern`` s (:meth:`add_task_restart_patterns`) on
+        the source networks are not carried over; set them on the merged
+        network afterward if needed.
+
+        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. See
+            :meth:`AlchemiscaleClient.set_network_state` for valid states.
+            Defaults to ``"active"``.
+        visualize
+            If ``True``, show 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 copy_network(
+        self,
+        network: ScopedKey,
+        scope: Scope,
+        name: str | None = None,
+        state: NetworkStateEnum | str = NetworkStateEnum.active,
+        visualize: bool = True,
+    ) -> ScopedKey:
+        """Copy an AlchemicalNetwork to a new Scope, carrying over every
+        existing ``complete`` Task and its associated
+        ProtocolDAGResultRefs.
+
+        Any ``Strategy`` (:meth:`set_network_strategy`) and
+        ``TaskRestartPattern`` s (:meth:`add_task_restart_patterns`) on
+        the source network are not carried over; set them on the copy
+        afterward if needed.
+
+        Parameters
+        ----------
+        network
+            The ScopedKey of the AlchemicalNetwork to copy.
+        scope
+            The destination Scope for the copy. This must be a *specific*
+            Scope; it must not contain wildcards.
+        name
+            Optional new name for the copy. If ``None``, the source
+            AlchemicalNetwork's name is preserved and the resulting copy
+            has the same gufe key as the source. If a different name is
+            given, the resulting copy gets a fresh gufe key derived from
+            the renamed content.
+        state
+            The starting state of the new AlchemicalNetwork. See
+            :meth:`AlchemiscaleClient.set_network_state` for valid states.
+            Defaults to ``"active"``.
+        visualize
+            If ``True``, show progress indicator.
+
+        Returns
+        -------
+        ScopedKey
+            The ScopedKey of the new, copied AlchemicalNetwork.
+        """
+        if not scope.specific():
+            raise ValueError(
+                f"`scope` '{scope}' contains wildcards ('*'); `scope` must be *specific*"
+            )
+
+        if not isinstance(network, ScopedKey):
+            network = ScopedKey.from_str(network)
+        if network.qualname != "AlchemicalNetwork":
+            raise ValueError(
+                f"ScopedKey '{network}' does not refer to an AlchemicalNetwork"
+            )
+
+        state = NetworkStateEnum(state)
+
+        data = dict(
+            scope=scope.to_dict(),
+            name=name,
+            state=state.value,
+        )
+
+        def post():
+            return self._post_resource(f"/networks/{network}/copy", data)
+
+        if visualize:
+            from rich.progress import Progress
+
+            with Progress(*self._rich_waiting_columns(), transient=False) as progress:
+                task = progress.add_task(
+                    f"Copying [bold]'{network}'[/bold] into scope "
+                    f"[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 merge_scopes(
+        self,
+        scopes: list[Scope],
+        target_scope: Scope,
+        visualize: bool = True,
+    ) -> list[ScopedKey]:
+        """Copy every AlchemicalNetwork in each of the given source Scopes
+        into ``target_scope``.
+
+        Each source AlchemicalNetwork is copied via :meth:`copy_network`,
+        preserving its name, its ``complete`` Tasks, and its associated
+        ProtocolDAGResultRefs.
+
+        Parameters
+        ----------
+        scopes
+            The source Scopes to copy from. Each must be a *specific*
+            Scope (no wildcards). The caller must have access to each.
+        target_scope
+            The destination Scope. Must be a *specific* Scope.
+        visualize
+            If ``True``, show a progress bar tracking the number of
+            networks copied.
+
+        Returns
+        -------
+        list[ScopedKey]
+            The ScopedKeys of all newly-copied AlchemicalNetworks in
+            ``target_scope``, in the order they were copied.
+        """
+        if not target_scope.specific():
+            raise ValueError(
+                f"`target_scope` '{target_scope}' contains wildcards ('*'); "
+                "`target_scope` must be *specific*"
+            )
+        if not scopes:
+            raise ValueError("`scopes` must contain at least one Scope")
+        for sc in scopes:
+            if not sc.specific():
+                raise ValueError(
+                    f"source scope '{sc}' contains wildcards ('*'); each source "
+                    "scope must be *specific*"
+                )
+
+        # gather every AlchemicalNetwork across the source scopes up front,
+        # regardless of state, so the progress bar has an accurate total
+        network_sks: list[ScopedKey] = []
+        for sc in scopes:
+            network_sks.extend(self.query_networks(scope=sc, state=None))
+
+        if visualize:
+            from rich.progress import Progress
+
+            with Progress(*self._rich_waiting_columns(), transient=False) as progress:
+                task = progress.add_task(
+                    f"Copying [bold]{len(network_sks)}[/bold] networks into "
+                    f"[bold]'{target_scope}'[/bold]...",
+                    total=len(network_sks),
+                )
+                new_sks: list[ScopedKey] = []
+                for sk in network_sks:
+                    new_sks.append(self.copy_network(sk, target_scope, visualize=False))
+                    progress.update(task, advance=1)
+        else:
+            new_sks = [
+                self.copy_network(sk, target_scope, visualize=False)
+                for sk in network_sks
+            ]
+
+        return new_sks
+
     def set_network_state(
         self, network: ScopedKey, state: NetworkStateEnum | str
     ) -> ScopedKey | None: