From 45696d8a6df22b30832a71165e99e8cac6524741 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Joan=20H=C3=A9risson?= Date: Tue, 13 Oct 2026 08:52:21 +0200 Subject: [PATCH] Support aggregating RetroRules datasets within one release --- README.md | 34 ++++++++- rrparser/Args.py | 2 +- rrparser/Downloads.py | 65 +++++++++++++++- rrparser/Parser.py | 112 ++++++++++++++++++++------- tests/test_aggregation.py | 154 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 338 insertions(+), 29 deletions(-) create mode 100644 tests/test_aggregation.py diff --git a/README.md b/README.md index cc52744..d3aa1da 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,10 @@ python -m rrparser --release 3.1.0 --dataset rhea --diameters 0-8 -o rules.csv # Equivalent selection using radii. python -m rrparser --dataset rhea --radii 0,1,2,3,4 -o rules.csv +# Aggregate datasets within one release (comma or + separator). +python -m rrparser --release 3.1.0 --dataset metanetx,rhea -d 0-8 -o rules.csv +python -m rrparser --release 3.1.0 --dataset all -o rules.csv + # Historical rr02 rules, restricted to retro usage and diameter 10. python -m rrparser --release rr02 --rule-type retro -d 10 -o rules.csv @@ -59,6 +63,26 @@ reused from cache, including offline. To update, list available releases and choose a version explicitly, or delete `retrorules-latest-.txt` from the cache to resolve `latest` again. The default remains pinned to `3.1.0`. +### Aggregate datasets + +Use `--dataset metanetx,rhea` or `--dataset metanetx+rhea` to combine datasets, +or `--dataset all` to include every dataset discovered for the selected release. +Repeated names are ignored. With rr02, `all` selects its only dataset, MetaNetX. + +Aggregation keeps one row per template ID. Shared templates combine their +source reactions, EC annotations, radii, and dataset provenance as unique +semicolon-separated values. Counts and radius bounds are recalculated, the +highest available score is retained, and validation must hold in every source. +Different SMARTS for the same template ID cause an error. Diameter/radius and +EC filtering, followed by optional score replacement, apply to the merged data. + +Each original dataset uses its own cache file; no combined archive is created. +`latest` resolves a single release containing all explicitly requested datasets, +so versions are never mixed. `all` and `latest` selections are recorded in +`retrorules-selection--.txt` for offline reuse. Delete that +selection file to discover the choices again. An explicit cached combination +such as `metanetx,rhea` also works offline without a selection file. + ### Diameter and radius selection `--diameters` accepts comma-separated values, inclusive ranges, and comparisons: @@ -130,7 +154,7 @@ archive cache. | --- | --- | | `--rules_file` | Local input file, or `retrorules` (default) to use a cached/downloaded dataset. | | `--release` | Version, `rr02`, or `latest`; default `3.1.0`. | -| `--dataset` | Dataset name; default `metanetx`. See `--list-releases`. | +| `--dataset` | Dataset name, comma/plus-separated names, or `all`; default `metanetx`. See `--list-releases`. | | `--list-releases` | Print available release/dataset pairs and exit. | | `--cache-dir`, `--rules-dir` | Source cache directory; default package-local `rrparser/cache`. | | `-d`, `--diameters` | List, range, or comparison; default all even diameters 0–20. | @@ -172,6 +196,14 @@ csv_text = parse_rules( radii="0,1,2", ) +# Aggregate datasets; dataset="all" is also supported. +parse_rules( + outfile="combined.csv", + release="3.1.0", + dataset="metanetx+rhea", + diameters="<=8", +) + # Parse a local file with EC filtering and replacement scores. parse_rules( outfile="filtered.csv", diff --git a/rrparser/Args.py b/rrparser/Args.py index 004934e..3a719f3 100644 --- a/rrparser/Args.py +++ b/rrparser/Args.py @@ -59,7 +59,7 @@ def _add_arguments(parser): parser.add_argument( "--dataset", default="metanetx", - help="RetroRules dataset (default: metanetx; see --list-releases)", + help="Dataset names (metanetx,rhea or metanetx+rhea), or all (default: metanetx; see --list-releases)", ) parser.add_argument( "-o", diff --git a/rrparser/Downloads.py b/rrparser/Downloads.py index ee02f2a..9831d44 100644 --- a/rrparser/Downloads.py +++ b/rrparser/Downloads.py @@ -87,7 +87,9 @@ def _selection_error(message, datasets=None): return ValueError(f"{message}\nAvailable releases and datasets:\n{choices}") -def fetch_dataset(rules_dir, release=DEFAULT_RELEASE, dataset="metanetx"): +def fetch_dataset( + rules_dir, release=DEFAULT_RELEASE, dataset="metanetx", available=None +): """Reuse cached rules; resolve latest only on the first download.""" if release == "rr02": if dataset != "metanetx": @@ -111,7 +113,7 @@ def fetch_dataset(rules_dir, release=DEFAULT_RELEASE, dataset="metanetx"): filename = cache / (f"retrorules-{release}-{dataset}.csv.gz") if release != "latest" and filename.is_file(): return str(filename) - datasets = list_retrorules() + datasets = available if available is not None else list_retrorules() if release == "latest": release = next( ( @@ -132,3 +134,62 @@ def fetch_dataset(rules_dir, release=DEFAULT_RELEASE, dataset="metanetx"): if latest: selection.write_text(release) return str(filename) + + +def fetch_datasets(rules_dir, release=DEFAULT_RELEASE, dataset="metanetx"): + """Select datasets from one release, preserving their individual caches.""" + names = sorted({name.strip() for name in dataset.replace("+", ",").split(",")}) + if any(not fullmatch(r"[a-z0-9_-]+", name) for name in names): + raise ValueError("Use dataset names separated by commas, or all") + if "all" in names and names != ["all"]: + raise ValueError("Use all alone, or a list of dataset names") + if release == "rr02" and names == ["all"]: + names = ["metanetx"] + if len(names) == 1 and names != ["all"]: + return [fetch_dataset(rules_dir, release, names[0])] + if release == "rr02": + raise ValueError("RetroRules rr02 is only available for metanetx") + if release != "latest" and not fullmatch(r"\d+\.\d+(?:\.\d+)?", release): + raise _selection_error(f"Invalid release {release!r}") + cache = Path(rules_dir) + selection = cache / f"retrorules-selection-{release}-{'-'.join(names)}.txt" + if selection.is_file(): + saved = selection.read_text().splitlines() + if len(saved) == 2 and fullmatch(r"\d+\.\d+(?:\.\d+)?", saved[0]): + selected = saved[1].split(",") + if selected and all(fullmatch(r"[a-z0-9_-]+", name) for name in selected): + files = [ + cache / f"retrorules-{saved[0]}-{name}.csv.gz" for name in selected + ] + if (release == "latest" or release == saved[0]) and all( + file.is_file() for file in files + ): + return list(map(str, files)) + if release != "latest" and names != ["all"]: + files = [cache / f"retrorules-{release}-{name}.csv.gz" for name in names] + if all(file.is_file() for file in files): + return list(map(str, files)) + available = list_retrorules() + versions = list( + dict.fromkeys(version for version, _ in available if version != "rr02") + ) + version = next( + ( + version + for version in versions + if (release == "latest" or release == version) + and ( + names == ["all"] or all((version, name) in available for name in names) + ) + ), + None, + ) + if version is None: + raise _selection_error( + f"Release/datasets unavailable: {release} / {dataset}", available + ) + if names == ["all"]: + names = [name for candidate, name in available if candidate == version] + files = [fetch_dataset(cache, version, name, available) for name in names] + selection.write_text(version + "\n" + ",".join(names)) + return files diff --git a/rrparser/Parser.py b/rrparser/Parser.py index f415ba7..d92ab98 100644 --- a/rrparser/Parser.py +++ b/rrparser/Parser.py @@ -10,10 +10,10 @@ from csv import QUOTE_ALL, QUOTE_NONE from logging import Logger, getLogger -from re import fullmatch +from re import fullmatch, split from tarfile import open as open_tar -from pandas import DataFrame, read_csv +from pandas import DataFrame, concat, read_csv from .Args import ( DEFAULT_DIAMETERS, @@ -21,7 +21,7 @@ DEFAULT_RULES_DIR, DEFAULT_RULES_FILE, ) -from .Downloads import RR02_FILE, fetch_dataset +from .Downloads import RR02_FILE, fetch_dataset, fetch_datasets LOGGER = getLogger(__name__) @@ -48,6 +48,84 @@ def read_ecnumbers(ec_file: str, logger: Logger = LOGGER) -> list[str]: return ec +def _read_rules(rules_file, **read_options): + if str(rules_file).endswith(".tar.gz"): + with open_tar(rules_file, mode="r:gz") as archive: + member = next( + ( + item + for item in archive + if item.name in (RR02_FILE, "./" + RR02_FILE) and item.isfile() + ), + None, + ) + if member is None: + raise ValueError( + f"RetroRules rr02 archive does not contain {RR02_FILE}" + ) + with archive.extractfile(member) as source: + return read_csv(source, **read_options) + else: + return read_csv(rules_file, **read_options) + + +def _merge_templates(frames): + """Union shared template annotations and retain the best source score.""" + templates = concat(frames, ignore_index=True) + duplicate = templates.duplicated("TEMPLATE_ID", keep=False) + if not duplicate.any(): + return templates + groups = templates[duplicate].groupby("TEMPLATE_ID", sort=False) + if (groups["TEMPLATE"].nunique() > 1).any(): + raise ValueError("Conflicting SMARTS for the same template ID across datasets") + + def union(values): + return ";".join( + sorted( + { + item + for value in values.dropna().astype(str) + for item in split(r"[;,]", value) + if item + } + ) + ) + + aggregations = { + column: "first" for column in templates.columns if column != "TEMPLATE_ID" + } + for column in ("REACTIONS", "ECS", "DATASETS", "RADII"): + if column in aggregations: + aggregations[column] = union + aggregations["SCORE"] = "max" + if "VALID" in aggregations: + aggregations["VALID"] = "all" + merged = groups.agg(aggregations).reset_index() + for column in ("REACTIONS", "ECS"): + if column + "_COUNT" in merged.columns: + merged[column + "_COUNT"] = merged[column].map( + lambda value: len(value.split(";")) if value else 0 + ) + if "RADII" in merged.columns: + radii = merged["RADII"].map( + lambda value: sorted(map(int, value.split(";"))) if value else [] + ) + merged["RADII"] = radii.map(lambda values: ";".join(map(str, values))) + for column, position in (("RADIUS_MIN", 0), ("RADIUS_MAX", -1)): + if column in merged.columns: + merged[column] = radii.map( + lambda values, position=position: ( + values[position] if values else None + ) + ) + result = concat([templates[~duplicate], merged], ignore_index=True) + return ( + result.set_index("TEMPLATE_ID") + .loc[templates["TEMPLATE_ID"].drop_duplicates()] + .reset_index() + ) + + def parse_rules( outfile: str | None, rules_file: str = DEFAULT_RULES_FILE, @@ -94,7 +172,7 @@ def parse_rules( release: str RetroRules version (default 3.1.0), rr02, or latest. dataset: str - Dataset discovered on retrorules.org (default metanetx). + Dataset name, comma/plus-separated names, or all (default metanetx). radii: str Comma-separated radii 0–10, as an alternative to diameters. @@ -116,36 +194,20 @@ def parse_rules( rule_type, diameters, input_format, output_format ) - # If 'rules_file' is set to RetroRules, then fetch RetroRules on Internet + rules_files = [rules_file] if rules_file == "retrorules": - rules_file = fetch_retrorules(rules_dir, logger, release, dataset) + rules_files = fetch_datasets(rules_dir, release, dataset) sep = "," logger.debug("Reading values...") # Read 'csv' as 'tsv' by specying separator read_options = { "sep": sep, - "dtype": {"RADII": str}, + "dtype": {"RADII": str, "ECS": str, "EC number": str}, "float_precision": "round_trip", } - if str(rules_file).endswith(".tar.gz"): - with open_tar(rules_file, mode="r:gz") as archive: - member = next( - ( - item - for item in archive - if item.name in (RR02_FILE, "./" + RR02_FILE) and item.isfile() - ), - None, - ) - if member is None: - raise ValueError( - f"RetroRules rr02 archive does not contain {RR02_FILE}" - ) - with archive.extractfile(member) as source: - rf = read_csv(source, **read_options) - else: - rf = read_csv(rules_file, **read_options) + frames = [_read_rules(filename, **read_options) for filename in rules_files] + rf = _merge_templates(frames) if len(frames) > 1 else frames[0] # Filter rules according to 'rule_type' and 'diameters' templates = "TEMPLATE_ID" in rf.columns diff --git a/tests/test_aggregation.py b/tests/test_aggregation.py new file mode 100644 index 0000000..628583c --- /dev/null +++ b/tests/test_aggregation.py @@ -0,0 +1,154 @@ +import gzip +from io import BytesIO, StringIO +from unittest.mock import patch + +import pandas as pd +import pytest + +from rrparser import build_args_parser, parse_rules +from rrparser.Downloads import fetch_datasets + + +def cache_dataset(cache, version, name, rows): + frame = pd.DataFrame( + rows, columns=["TEMPLATE_ID", "TEMPLATE", "RADII", "SCORE", "ECS", "REACTIONS"] + ) + frame["DATASETS"] = name + frame["VALID"] = True + for column in ("REACTIONS", "ECS"): + frame[column + "_COUNT"] = frame[column].map( + lambda value: len(value.split(";")) if value else 0 + ) + frame["RADIUS_MIN"] = frame["RADII"].map( + lambda value: min(map(int, value.split(";"))) + ) + frame["RADIUS_MAX"] = frame["RADII"].map( + lambda value: max(map(int, value.split(";"))) + ) + filename = cache / f"retrorules-{version}-{name}.csv.gz" + filename.write_bytes(gzip.compress(frame.to_csv(index=False).encode())) + return filename + + +@pytest.mark.parametrize( + "selection", ["metanetx,rhea", "rhea+metanetx", "metanetx,rhea,metanetx"] +) +def test_cached_datasets_are_merged_before_filtering(tmp_path, selection): + sources = [ + cache_dataset( + tmp_path, + "3.1.0", + "metanetx", + [ + ("shared", "[C]>>[O]", "0;1", 0.2, "1.1", "MNXR1"), + ("mnx", "[N]>>[O]", "9", 0.5, "3.1", "MNXR2"), + ], + ), + cache_dataset( + tmp_path, + "3.1.0", + "rhea", + [ + ("shared", "[C]>>[O]", "1;10", 0.8, "2.2", "RHEA:1"), + ("rhea", "[O]>>[N]", "5", 0.4, "4.1", "RHEA:2"), + ], + ), + ] + originals = [source.read_bytes() for source in sources] + args = build_args_parser().parse_args(["--dataset", selection]) + with patch( + "rrparser.Downloads.urlopen", side_effect=AssertionError("Network used") + ): + output = parse_rules(None, rules_dir=str(tmp_path), dataset=args.dataset) + filtered = parse_rules( + None, rules_dir=str(tmp_path), dataset=selection, radii="10", ec=["1.1"] + ) + frame = pd.read_csv(StringIO(output)) + assert set(frame["Rule ID"]) == {"shared", "mnx", "rhea"} + shared = frame.set_index("Rule ID").loc["shared"] + assert shared["RADII"] == "0;1;10" + assert shared["Score"] == 0.8 + assert shared["EC number"] == "1.1;2.2" + assert shared["REACTIONS"] == "MNXR1;RHEA:1" + assert shared["DATASETS"] == "metanetx;rhea" + assert shared["ECS_COUNT"] == shared["REACTIONS_COUNT"] == 2 + assert shared["RADIUS_MIN"] == 0 + assert shared["RADIUS_MAX"] == 10 + assert pd.read_csv(StringIO(filtered))["Rule ID"].tolist() == ["shared"] + assert [source.read_bytes() for source in sources] == originals + + +def test_all_is_discovered_and_can_be_reused_offline(tmp_path): + names = ["metanetx", "rhea", "uspto"] + for name in names: + cache_dataset(tmp_path, "3.1.0", name, [(name, "[C]>>[O]", "0", 0.5, "", name)]) + page = "".join( + f'CSV' for name in names + ).encode() + with patch("rrparser.Downloads.urlopen", return_value=BytesIO(page)) as request: + result = parse_rules(None, rules_dir=str(tmp_path), dataset="all") + request.assert_called_once() + assert set(pd.read_csv(StringIO(result))["Rule ID"]) == set(names) + with patch( + "rrparser.Downloads.urlopen", side_effect=AssertionError("Network used") + ): + assert parse_rules(None, rules_dir=str(tmp_path), dataset="all") == result + + +def test_latest_aggregation_uses_one_common_release(tmp_path): + for name in ["metanetx", "rhea"]: + cache_dataset(tmp_path, "3.0", name, [(name, "[C]>>[O]", "0", 0.5, "", name)]) + page = b"""CSV + CSV + CSV""" + with patch("rrparser.Downloads.urlopen", return_value=BytesIO(page)): + files = fetch_datasets(tmp_path, "latest", "metanetx,rhea") + assert all("3.0" in filename for filename in files) + with patch( + "rrparser.Downloads.urlopen", side_effect=AssertionError("Network used") + ): + assert fetch_datasets(tmp_path, "latest", "rhea+metanetx") == files + + +@pytest.mark.parametrize("selection", ["all,rhea", "metanetx,", "../rhea"]) +def test_invalid_selections(selection, tmp_path): + with pytest.raises(ValueError): + fetch_datasets(tmp_path, dataset=selection) + + +def test_conflicting_templates_are_rejected(tmp_path): + for name, template in [("metanetx", "[C]>>[O]"), ("rhea", "[O]>>[C]")]: + cache_dataset(tmp_path, "3.1.0", name, [("same", template, "1", 1, "", name)]) + with pytest.raises(ValueError, match="Conflicting SMARTS"): + parse_rules(None, rules_dir=str(tmp_path), dataset="metanetx,rhea") + + +def test_unknown_dataset_is_rejected_before_downloading(tmp_path): + page = b'CSV' + with patch("rrparser.Downloads.urlopen", return_value=BytesIO(page)) as request: + with pytest.raises(ValueError, match="Available releases"): + fetch_datasets(tmp_path, "3.1.0", "rhea,unknown") + request.assert_called_once() + assert not list(tmp_path.iterdir()) + + +def test_merged_score_replacement_and_missing_annotations(tmp_path): + cache_dataset( + tmp_path, "3.1.0", "rhea", [("shared", "[C]>>[O]", "0", 0.8, "1.1", "RHEA:1")] + ) + cache_dataset( + tmp_path, "3.1.0", "uspto", [("shared", "[C]>>[O]", "10", None, "", "USPTO:1")] + ) + scores = tmp_path / "scores.csv" + scores.write_text("Rule ID,Score\nshared,0.9\n") + result = parse_rules( + None, + rules_dir=str(tmp_path), + dataset="rhea,uspto", + radii="10", + scores_file=str(scores), + ) + shared = pd.read_csv(StringIO(result)).iloc[0] + assert shared["Score"] == 0.9 + assert shared["DATASETS"] == "rhea;uspto" + assert shared["EC number"] == 1.1