diff --git a/openfe_benchmarks/scripts/_tmp_alchemiscale_gather.py b/openfe_benchmarks/scripts/_tmp_alchemiscale_gather.py new file mode 100644 index 0000000..82a2689 --- /dev/null +++ b/openfe_benchmarks/scripts/_tmp_alchemiscale_gather.py @@ -0,0 +1,141 @@ +import click +import pathlib +from gufe.archival import AlchemicalArchive +from gufe import AlchemicalNetwork +from alchemiscale import AlchemiscaleClient, ScopedKey, Scope +import logging +import bz2 + +logger = logging.getLogger(__name__) + +def _configure_example_logging(level=logging.INFO): + handler = logging.StreamHandler() + handler.setFormatter( + logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s") + ) + # Attach to this module's logger so output appears when running the example + logger.addHandler(handler) + logger.setLevel(level) + # Optionally enable package-wide logs: + logging.getLogger("openfe_benchmarks").setLevel(level) + +@click.command() +@click.option("--network", type=click.Path(dir_okay=False, file_okay=True, exists=True, path_type=pathlib.Path), help="Path to the submitted alchemical network with the scope key saved as into the name field for which an archive should be made.", default=None) +@click.option("--network-key", type=str, help="The scope key of the network for which an archive should be made.", default=None) +@click.option("--org", type=str, help="Organization name used to build the scope to check can not be used with --network/network-key.", default=None) +@click.option("--campaign", type=str, help="Campaign name used to build the scope to check can not be used with --network/network-key.", default=None) +@click.option("--project", type=str, help="Project name used to build the scope to check can not be used with --network/network-key.", default=None) +@click.option("--allow-partial/--no-allow-partial", is_flag=True, default=False, help="Allow archive creation for an incomplete network.") +@click.option("--output", type=click.Path(dir_okay=True, path_type=pathlib.Path), help="Path to the output directory where the archive(s) should be made.", required=True) +def main(network, network_key, org, campaign, project, allow_partial, output): + """ + Create an AlchemicalArchive for a given alchemical network which has been submitted to Alchemiscale. + The network can be specified either by providing the path to the submitted network JSON file, the scoped key or by specifying the + scope (org, campaign, project). Archives for all networks + found within that scope will be created. + + Parameters + ---------- + network : pathlib.Path optional + The path to the submitted alchemicalnetwork JSON file which should have the network key set as the name field on the object + network_key : str optional + The alchemical network scope key + org : str optional + The organization name used to build the scope to check can not be used with --network/network-key. + campaign : str optional + The campaign name used to build the scope to check can not be used with --network/network-key. + project : str optional + The project name used to build the scope to check can not be used with --network/network-key. + allow_partial : bool + Whether to allow archive creation for an incomplete network, default is False which means only networks with all edges completed will be archived. + output : pathlib.Path + Path to the output directory where the archive(s) should be made. + + Notes + ----- + - The archive will be saved to the output folder with the same name as the network key in compressed JSON format. + + Examples + -------- + - Create an archive for a specific network specified by scope + ``` + python _tmp_alchemiscale_gather.py --org openfe --campaign my_first_run --project tyk2 --output local_archives + ``` + - Create an archive for each of the networks in a given campaign + ``` + python _tmp_alchemiscale_gather.py --org openfe --campaign my_first_run --output local_archives + ``` + - Create an archive for a specific scoped key + ``` + python _tmp_alchemiscale_gather.py --network-key my_network_key --output local_archives + ``` + """ + # we can not use scope with network/network-key + if network_key is not None and (network is not None or org is not None or campaign is not None or project is not None): + raise ValueError( + "Scope searching (org, campaign, project) and --network can not be used when --network-key is provided these are mutually exclusive.") + if network is not None and (org is not None or campaign is not None or project is not None): + raise ValueError( + "Scope searching (org, campaign, project) can not be used when --network is provided these are mutually exclusive.") + + output.mkdir(exist_ok=True, parents=True) + _configure_example_logging() + client = AlchemiscaleClient(api_url="https://api.alchemiscale.org") + + networks_by_keys = {} + # load the local network if provided + if network is not None: + logger.info(f"Loading local network from {network}") + alchem_network = AlchemicalNetwork.from_json(network.as_posix()) + alchem_network_key = ScopedKey.from_str(alchem_network.name) + networks_by_keys[alchem_network_key] = alchem_network + else: + if network_key is not None: + logger.info(f"Loading network from key: {network_key}") + alchem_network_key = ScopedKey.from_str(network_key) + alchem_network = client.get_network(alchem_network_key) + networks_by_keys[alchem_network_key] = alchem_network + else: + # use the scope to query for network keys + query_scope = Scope(org=org, campaign=campaign, project=project) + logger.info(f"Querying Alchemiscale for submissions with scope: {query_scope}") + network_keys = client.query_networks(scope=query_scope) + + logger.info(f"Number of networks found: {len(network_keys)}") + for network_key in network_keys: + logger.info(f"Loading network with key: {network_key}") + network = client.get_network(network_key, visualize=False) + networks_by_keys[network_key] = network + + logger.info(f"Creating archives") + for network_key, network in networks_by_keys.items(): + logger.info(f"Creating archive for network: {network_key}") + # first check the status + status = client.get_network_status(network_key) + status.pop("complete") + if any(v == 0 for v in status.values()): + if not allow_partial: + raise ValueError(f"Network {network_key} is incomplete to archive run with --allow_partial.") + else: + logger.warning(f"Network {network_key} has incomplete edges but will be archived anyway as --allow-partial is set.") + + + # we need the transformations and the results to build the archive + transformation_results = [] + alchemiscale_network_results = client.get_network_results(network_key, return_protocoldagresults=True).items() + for transformation_key, raw_results in alchemiscale_network_results: + transformation = client.get_transformation(transformation_key, visualize=False) + transformation_results.append((transformation, raw_results)) + + archive = AlchemicalArchive( + network=network, + transformation_results=transformation_results, + ) + # write the archive to a compressed format + with bz2.open(output / f"{network_key}.json.bz2", "wb") as outfile: + outfile.write(archive.to_json().encode()) + + logger.info(f"Archive for network {network_key} created and saved to {output / f'{network_key}.json.bz2'}") + +if __name__ == "__main__": + main()