diff --git a/chebILP/cli.py b/chebILP/cli.py index f5cfde8..b8317e3 100644 --- a/chebILP/cli.py +++ b/chebILP/cli.py @@ -484,7 +484,7 @@ def build_parser() -> argparse.ArgumentParser: sp_learn.add_argument("--selection_k", type=int, default=10, help="Number of predicates selection with selection_mode (required if selection_mode is set).") sp_learn.add_argument("--max_vars", type=int, default=6, help="Maximum number of variables in learned rules.") sp_learn.add_argument("--max_body", type=int, default=8, help="Maximum number of body literals in learned rules.") - sp_learn.add_argument("--max_clauses", type=int, default=2, help="Maximum number of clauses in the learned program.") + sp_learn.add_argument("--max_clauses", type=int, default=2, help="Maximum number of clauses in the learned program. Inert: Popper forces max_rules=1 unless recursion or predicate invention is enabled, and the noisy MDL combiner caps nothing.") sp_learn.add_argument("--mdl_weight_fn", type=int, default=1, help="Weight β for false negatives in MDL cost (default: 1).") sp_learn.add_argument("--mdl_weight_fp", type=int, default=1, help="Weight γ for false positives in MDL cost (default: 1).") sp_learn.add_argument("--mdl_weight_size", type=int, default=1, help="Weight α for program size in MDL cost (default: 1).") diff --git a/chebILP/evaluation/clingo_eval.py b/chebILP/evaluation/clingo_eval.py index a568dac..58b8a8f 100644 --- a/chebILP/evaluation/clingo_eval.py +++ b/chebILP/evaluation/clingo_eval.py @@ -1,6 +1,33 @@ +import os +import re + +from tqdm import tqdm + from chebILP.utils import split_prolog_literals +def _patch_clingo_logger_decode() -> None: + """Decode clingo logger messages with ``errors="replace"``. + + clingo's logger callback strict-decodes each message as UTF-8, and its + ``onerror`` handler hard-exits the process (``os._exit``) on failure — so a + single non-UTF-8 byte in a warning/info message kills the whole run and no + ``try/except`` can catch it. A custom ``logger=`` does not help: the decode + runs before the handler is called. Patching ``clingo.core._to_str`` makes the + decode tolerant so grounding messages can never take the process down. + """ + from clingo import core + + if getattr(core, "_chebilp_safe_to_str", False): + return + + def _safe_to_str(c_str) -> str: + return core._ffi.string(c_str).decode(errors="replace") + + core._to_str = _safe_to_str + core._chebilp_safe_to_str = True + + def filter_impossible_rules(rules: list[str], predicates_in_bk: list[str]): # for every predicate name in the rule body, check if it exists in background_facts predicates_in_bk = [p[0] for p in predicates_in_bk] @@ -14,16 +41,20 @@ def filter_impossible_rules(rules: list[str], predicates_in_bk: list[str]): print(f"Filtered out {len(rules) - len(rules_filtered)} impossible rules. Remaining rules: {len(rules_filtered)}") return rules_filtered -def ground_extensions(rules: list[str], background_facts: list[str], target_labels: list[str], timeout: float|None=None) -> dict[str, list[tuple[str, ...]]]: +def ground_extensions(rules: list[str], background_facts: list[str], target_labels: list[str], timeout: float|None=None, message_sink: list|None=None) -> dict[str, list[tuple[str, ...]]]: """Ground ``rules`` over ``background_facts``; return each target's derived argument tuples. The target predicates may have any arity — the arguments are returned as they were derived and it is up to the caller to interpret them. Note that clingo treats ``p/1`` and ``p/2`` as different predicates, so one name can yield tuples of differing width. + + Pass ``message_sink`` to collect clingo's diagnostics into it instead of letting them go + to stderr — the caller can then report them once rather than once per grounding. """ import clingo - ctl = clingo.Control() + _patch_clingo_logger_decode() + ctl = clingo.Control(logger=(lambda code, msg: message_sink.append(msg)) if message_sink is not None else None) ctl.add("base", [], "\n".join(background_facts)) try: ctl.add("base", [], "\n".join(rules)) @@ -62,6 +93,199 @@ def _collect(model): return extensions +# How much address space a grounding child may claim ON TOP OF what it already inherited at +# fork. It is a budget, not an absolute ``RLIMIT_AS``: the auxiliary-rule generator loads +# sentence-transformers/torch before it ever grounds, and torch's reserved VA arenas put the +# parent's VSZ around 42 GiB while its RSS stays near 2 GiB. An absolute cap below that is +# worse than none — ``setrlimit`` accepts it, then every allocation in the child fails, so +# even a program that grounds in a millisecond comes back as a memory failure. +GROUNDING_MEMORY_BUDGET_BYTES = 4 * 1024**3 + +# Molecules per clingo instance. A clause whose variables are not all joined to the molecule +# grounds as a cross product over EVERY atom in the instance, so the cost of one bad rule is +# quadratic in the number of molecules ground together. Grounding a whole class at once (650+ +# molecules, ~13k ring atoms) puts that at ~10^8 ground atoms; batching caps it near 10^6. +# For a well-formed program — every variable joined to its molecule — no clause can span two +# molecules, so the batch size does not affect the result at all. +GROUNDING_BATCH_SIZE = 25 + + +def _address_space_limit(budget: int) -> int | None: + """``RLIMIT_AS`` value leaving ``budget`` bytes of headroom above what is mapped now.""" + try: + with open("/proc/self/statm") as f: + mapped_pages = int(f.read().split()[0]) + except (OSError, ValueError, IndexError): + return None + return mapped_pages * os.sysconf("SC_PAGE_SIZE") + budget + + +# "atom does not occur in any rule head" spans two lines, the atom on the second. +_UNDEFINED_ATOM_RE = re.compile( + r"atom does not occur in any rule head:\s*([a-z_][A-Za-z0-9_]*)\s*(\(([^)]*)\))?", re.IGNORECASE +) + +# ":5093:17-18: info: " — clingo's severity, as it spells it in the message. +_INFO_MESSAGE_RE = re.compile(r":\s*info:", re.IGNORECASE) + + +def _summarize_clingo_messages(per_group: list[list[str]]) -> None: + """Print one line for what clingo said, however many groundings it said it in. + + Every batch re-reports the same diagnostics, so the raw stream is a handful of messages + times the batch count. Undefined predicates are the signal worth keeping — a rule + referencing something nothing defines has an empty body — but "undefined" is reported + *per grounding*, and a batch whose molecules happen to contain no sulfur reports ``s/1`` + exactly like a genuine typo. For an atom-level predicate only absence from **every** batch + means anything, and with a single batch (one molecule, as when building a prediction + tensor) not even that — so those are reported only when there are batches to compare. + An undefined ``aux_`` name is never molecule variation: nothing but a rule program can + define one, so a missing library dependency is reported however few batches there were. + Warnings and errors are printed once, deduplicated. + + Clingo's other ``info:`` messages are dropped. They describe the *program* — a global + variable used as an aggregate tuple, an empty interval — so they are identical for every + molecule and every batch, and the ``:LINE`` they point at is an offset into a + concatenation that differs from call to call. One such clause therefore prints once per + molecule while naming a line nobody can look up. Program defects belong to the static + checks in ``predicate_generation.auxiliary_rules``, which report them once per class + against the real source file. + """ + seen_in: dict[str, set[int]] = {} + other: set[str] = set() + for index, messages in enumerate(per_group): + for message in messages: + m = _UNDEFINED_ATOM_RE.search(message) + if m is None: + if not _INFO_MESSAGE_RE.search(message): + other.add(message.strip()) + continue + arity = len(m.group(3).split(",")) if m.group(3) and m.group(3).strip() else 0 + seen_in.setdefault(f"{m.group(1)}/{arity}", set()).add(index) + + undefined = sorted( + name for name, groups in seen_in.items() + if len(groups) == len(per_group) and (name.startswith("aux_") or len(per_group) > 1) + ) + # tqdm.write instead of print: callers run this inside a progress bar, and a plain + # print leaves the bar's line behind. Without a bar it behaves like print. + if undefined: + tqdm.write(f" clingo: {len(undefined)} predicate(s) referenced but never defined " + f"(their rule bodies are empty): {', '.join(undefined)}") + for message in sorted(other): + tqdm.write(f" clingo: {message}") + + +def _ground_groups(rules, fact_groups, target_labels, timeout): + """Ground ``rules`` over each group of facts separately, merging the derived tuples.""" + merged: dict[str, list[tuple[str, ...]]] = {} + per_group: list[list[str]] = [] + for group in fact_groups: + messages: list = [] + per_group.append(messages) + derived = ground_extensions(rules, group, target_labels, timeout=timeout, message_sink=messages) + for name, tuples in derived.items(): + merged.setdefault(name, []).extend(tuples) + _summarize_clingo_messages(per_group) + return merged + + +def _ground_extensions_child(queue, rules, fact_groups, target_labels, timeout, memory_budget): + """Child-process body: cap address space, then ground and return the result via ``queue``.""" + try: + import resource + + limit = _address_space_limit(memory_budget) + if limit is not None: + resource.setrlimit(resource.RLIMIT_AS, (limit, limit)) + except (ImportError, ValueError, OSError): + pass # cap unavailable: the parent's wall-clock ceiling is then the only backstop + try: + queue.put(("ok", _ground_groups(rules, fact_groups, target_labels, timeout))) + except MemoryError as e: + # Reporting the failure allocates, and under an exhausted RLIMIT_AS that raises a + # second MemoryError which kills the child before it can say why — leaving only its + # traceback on the shared stderr. Clearing the traceback releases the grounding + # frames it pins, and with them clingo's Control and the partial result. + e.__traceback__ = None + queue.put(("err", "grounding exceeded the memory limit")) + except Exception as e: # report any failure back rather than dying silently + message = str(e).strip().splitlines()[0] if str(e).strip() else repr(e) + e.__traceback__ = None + queue.put(("err", message)) + + +def ground_extensions_isolated( + rules: list[str], + fact_groups: list[list[str]], + target_labels: list[str], + timeout: float | None = None, + total_timeout: float | None = None, + memory_budget: int = GROUNDING_MEMORY_BUDGET_BYTES, +) -> dict[str, list[tuple[str, ...]]]: + """Ground ``rules`` over each group in ``fact_groups``, in a memory-capped forked child. + + Each group is a separate clingo instance and the derived tuples are merged, so a clause + can never join facts from two different groups. Callers with a single flat fact list + pass ``[facts]``; :func:`derive_rule_extensions` splits by molecule instead, which is + what keeps an unjoined-variable clause from grounding as a class-wide cross product. + + A runaway grounding dies inside the child (memory cap or, as a last resort, the OOM + killer) instead of taking the whole process down uncatchably. On any child failure — + memory cap hit, killing signal, timeout, or grounding error — this raises + ``RuntimeError`` so the caller can reject the offending programs and continue. + + ``timeout`` is the per-group solve budget; ``total_timeout`` is the wall-clock ceiling + for the whole child, and defaults to ``timeout + 60`` for a single group. Pass one of + them: without any ceiling a child that wedges rather than dying leaves the parent + polling for ever. + + Falls back to an in-process call where forking is unavailable (e.g. Windows). + """ + import multiprocessing + import queue as queue_mod + import time + + try: + ctx = multiprocessing.get_context("fork") + except ValueError: + return _ground_groups(rules, fact_groups, target_labels, timeout) + + result_queue = ctx.Queue() + proc = ctx.Process( + target=_ground_extensions_child, + args=(result_queue, rules, fact_groups, target_labels, timeout, memory_budget), + ) + proc.start() + + # Wall-clock ceiling for the whole child, distinct from the per-group ``timeout``. + join_timeout = total_timeout + if join_timeout is None and timeout is not None: + join_timeout = timeout * max(len(fact_groups), 1) + 60 + deadline = None if join_timeout is None else time.monotonic() + join_timeout + + # Drain the queue while the child runs: a large result can exceed the pipe buffer, so + # the child blocks on flush until we read — joining first would deadlock. Polling also + # lets us notice a child that was killed (OOM/SIGKILL) without producing a result. + while True: + try: + status, payload = result_queue.get(timeout=0.5) + break + except queue_mod.Empty: + if not proc.is_alive(): + proc.join() + raise RuntimeError(f"grounding died before returning (exit code {proc.exitcode})") + if deadline is not None and time.monotonic() > deadline: + proc.terminate() + proc.join() + raise RuntimeError(f"grounding exceeded the {join_timeout}s time limit") + + proc.join() + if status == "ok": + return payload + raise RuntimeError(payload) + + def evaluate_with_clingo(rules: list[str], background_facts: list[str], target_labels: list[str], examples: list, predicates_in_bk: list[str]|None=None, timeout: float|None=None): """Which of ``examples`` each target predicate holds for, matching ``label(example)``. diff --git a/chebILP/evaluation/test.py b/chebILP/evaluation/test.py index 69c2415..59eb13c 100644 --- a/chebILP/evaluation/test.py +++ b/chebILP/evaluation/test.py @@ -23,120 +23,81 @@ def _silent_clingo_logger(code, message): _WORKER_STATE: dict = {} -def _qualify_aux_predicates(programs, aux_library_dir): - """Namespace auxiliary-predicate references per class so they cannot collide. - - Auxiliary predicate names (e.g. ``aux_long_aliphatic_chain``) are NOT unique - across classes: many classes independently define a predicate with the same - sanitized name but *different* logic. When every class's program is grounded - against one shared background knowledge, those names would collide and a single - (arbitrary) implementation would be used for all of them — silently producing the - wrong labels. To keep the shared-BK optimization sound, rewrite each program's aux - references to a class-qualified name (``_``, the class id being - unique) and emit that class's own extension under the same name. - - Returns ``(rewritten_programs, aux_specs)`` where ``rewritten_programs`` maps - class_id -> program text with qualified aux names, and ``aux_specs`` is a list of - ``(qualified_name, source_file)`` telling a worker which predicates to (re)load. +def _aux_predicate_specs(programs, aux_library_dir): + """Which auxiliary predicates (llm_generated_fgs) the learned programs reference. + + An ``aux_`` name identifies one implementation library-wide: ``add_program_to_library`` + disambiguates a name that would back different code with the class id, so a shared + background knowledge can hold every class's predicates side by side under their own + names. Only the ones some program actually mentions are worth computing. + + Returns a list of ``(name, source_file)`` telling a worker which predicates to load. """ import re from chebILP.predicate_generation.auxiliary_predicates import load_auxiliary_predicates - rewritten: dict = {} - specs: dict[str, str] = {} # qualified_name -> source_file + specs: dict[str, str] = {} for cls_id, prog in programs.items(): used = set(re.findall(r"\baux_\w+", prog)) if not used: - rewritten[cls_id] = prog continue - # Map this class's own sanitized aux names to their source files. - name_to_source = { - p.name: p.source_file - for p in load_auxiliary_predicates(cls_id, library_dir=aux_library_dir) - } - rename: dict[str, str] = {} - for name in used: - source_file = name_to_source.get(name) - if source_file is None: - continue # program references an aux this class does not define - qname = f"{name}_{cls_id}" - rename[name] = qname - specs.setdefault(qname, source_file) - rewritten[cls_id] = re.sub( - r"\baux_\w+", lambda m: rename.get(m.group(0), m.group(0)), prog - ) - return rewritten, list(specs.items()) + for p in load_auxiliary_predicates(cls_id, library_dir=aux_library_dir): + if p.name in used: + specs.setdefault(p.name, p.source_file) + return list(specs.items()) -def _qualify_aux_rules(programs, rule_library_dir): - """Namespace auxiliary-RULE references per class (llm_generated_rules). +def _collect_rule_programs(programs, rule_library_dir): + """Rule programs (llm_generated_rules) needed to evaluate the learned programs. - Like :func:`_qualify_aux_predicates`, but for ASP rule programs: the same sanitized - ``aux_`` name can back different rules in different classes, and every class's program - is grounded against one shared background knowledge. Rewrite each program's aux - references to a class-qualified name (``_``) and return the matching - qualified :class:`RuleProgram` objects so the background emits each class's own - extension under that name. + Returns ``(rule_programs, dependency_programs)``: the programs the classes chose, and + the library programs those build on. ``class_map.json`` records only the former, so + without the latter every layered rule grounds against an empty body and derives nothing. + Kept apart because only the chosen programs' extensions belong in the background — + the dependencies exist to make them derivable, exactly as in ``ILPProblemBuilder.build_bk``. - Returns ``(rewritten_programs, qualified_rule_programs)``. Rule programs carry only - strings, so (unlike Python aux predicates) they are picklable and shipped to workers - directly rather than reloaded from disk. + Names are library-wide unique (``add_rule_to_library`` disambiguates colliding content + with the class id), so the programs go into the shared background as they are. Rule + programs carry only strings, so (unlike Python aux predicates) they are picklable and + shipped to workers directly rather than reloaded from disk. - EVERY aux name a class's programs mention is renamed, not just the heads the learned - program uses, and the class's whole rule set is returned rather than the subset it uses. - Both are required because the programs are grounded together: a head the learned program - never mentions may still be the helper another head depends on, and an unqualified helper - would otherwise merge with a same-named helper from a different class. + A class's *whole* rule set is loaded, not just the heads its learned program uses: a head + the program never mentions may still be the helper another head depends on. """ import re - from chebILP.predicate_generation.auxiliary_rules import RuleProgram, load_class_rules + from chebILP.predicate_generation.auxiliary_rules import load_class_rules, resolve_rule_dependencies - rewritten: dict = {} - qualified: dict[str, RuleProgram] = {} + chosen = {} for cls_id, prog in programs.items(): if not re.search(r"\baux_\w+", prog): - rewritten[cls_id] = prog continue - rps = load_class_rules(cls_id, library_dir=rule_library_dir) - rename = { - name: f"{name}_{cls_id}" - for rp in rps - for name in re.findall(r"\baux_\w+", rp.source) - } - substitute = lambda text: re.sub(r"\baux_\w+", lambda m: rename.get(m.group(0), m.group(0)), text) - for rp in rps: - qname = rename[rp.name] - if qname not in qualified: - qualified[qname] = RuleProgram( - name=qname, - description=rp.description, - source=substitute(rp.source), - source_file=rp.source_file, - ) - rewritten[cls_id] = substitute(prog) - return rewritten, list(qualified.values()) - - -def _load_qualified_aux(specs): - """Reload auxiliary predicates for a worker under their class-qualified names. - - ``specs`` is a list of ``(qualified_name, source_file)`` produced by - ``_qualify_aux_predicates``. The predicates' ``extension`` functions are - ``exec``-compiled and therefore not picklable, so each worker rebuilds them from - disk rather than receiving them; the qualified name namespaces each class's - version so identically-named predicates from different classes cannot collide. + for rp in load_class_rules(cls_id, library_dir=rule_library_dir): + chosen.setdefault(rp.name, rp) + + rule_programs = list(chosen.values()) + dependencies = [ + dep for dep in resolve_rule_dependencies(rule_programs, rule_library_dir) + if dep.name not in chosen + ] + return rule_programs, dependencies + + +def _load_aux_predicates(specs): + """Reload auxiliary predicates for a worker from ``(name, source_file)`` specs. + + The predicates' ``extension`` functions are ``exec``-compiled and therefore not + picklable, so each worker rebuilds them from disk rather than receiving them. """ from chebILP.predicate_generation.auxiliary_predicates import load_program_source preds = [] - for qname, source_file in specs: + for _, source_file in specs: try: with open(source_file, "r", encoding="utf-8") as f: pred = load_program_source(f.read(), source_file=source_file) except OSError: pred = None if pred is not None: - pred.name = qname preds.append(pred) return preds @@ -159,16 +120,15 @@ def _quiet_aux_logging(): def _worker_init(state, aux_load_args): """Pool initializer: stash the shared config and (re)load auxiliary predicates. - ``aux_load_args`` is the ``aux_specs`` list from ``_qualify_aux_predicates`` - (``(qualified_name, source_file)`` pairs) or ``None``. The predicates' - ``extension`` functions are ``exec``-compiled and therefore not picklable, so each - worker reloads them from disk rather than receiving them. + ``aux_load_args`` is the spec list from ``_aux_predicate_specs`` (``(name, source_file)`` + pairs) or ``None``. The predicates' ``extension`` functions are ``exec``-compiled and + therefore not picklable, so each worker reloads them from disk rather than receiving them. """ global _WORKER_STATE _quiet_aux_logging() _WORKER_STATE = dict(state) _WORKER_STATE["aux_predicates"] = ( - _load_qualified_aux(aux_load_args) if aux_load_args is not None else None + _load_aux_predicates(aux_load_args) if aux_load_args is not None else None ) @@ -197,7 +157,8 @@ def _evaluate_molecule(payload): row_df, predicate_set=st["predicate_set"], aux_predicates=st["aux_predicates"], aux_timeout=st["aux_timeout"], aux_failures=aux_failures, fowl_smarts=st.get("fowl_smarts"), - rule_programs=st.get("rule_programs"), computed_facts=st.get("computed_facts", False), + rule_programs=st.get("rule_programs"), rule_dependencies=st.get("rule_dependencies"), + computed_facts=st.get("computed_facts", False), ) except Exception as e: # noqa: BLE001 return mol_id, [], f"error: {e}", aux_failures @@ -309,26 +270,24 @@ def build_ilp_preds_tensor( except RuntimeError as e: print(f" Skipping ChEBI:{cls_id} — failed to parse program: {e}") - # Auxiliary predicates (llm_generated_fgs) are class-specific: the same sanitized name - # can mean different things in different classes. Qualify each program's aux references - # (and the extensions we emit for them) per class so they cannot collide in the single - # shared background knowledge. The predicates are then (re)loaded once per worker (their + # Auxiliary predicates (llm_generated_fgs) are gathered across all classes and computed + # into the single shared background knowledge. They are (re)loaded once per worker (their # exec-compiled extension functions are not picklable); ``aux_load_args`` is the spec # list a worker needs to reload them. aux_load_args = None if predicate_set == "llm_generated_fgs": - valid_programs, aux_specs = _qualify_aux_predicates(valid_programs, aux_library_dir) - aux_load_args = aux_specs - print(f"{len(aux_specs)} distinct auxiliary predicate implementation(s) referenced " + aux_load_args = _aux_predicate_specs(valid_programs, aux_library_dir) + print(f"{len(aux_load_args)} distinct auxiliary predicate implementation(s) referenced " f"by programs") - # llm_generated_rules: qualify aux rule references per class and recompute their - # extensions in the background. RuleProgram objects are picklable, so they ride in - # worker_state directly (no per-worker reload like the Python predicates need). - rule_programs = None + # llm_generated_rules: gather the classes' rule programs (plus the library programs they + # build on) and recompute their extensions in the background. RuleProgram objects are + # picklable, so they ride in worker_state directly. + rule_programs = rule_dependencies = None if predicate_set == "llm_generated_rules": - valid_programs, rule_programs = _qualify_aux_rules(valid_programs, aux_library_dir) - print(f"{len(rule_programs)} distinct auxiliary rule(s) referenced by programs") + rule_programs, rule_dependencies = _collect_rule_programs(valid_programs, aux_library_dir) + print(f"{len(rule_programs)} distinct auxiliary rule(s) referenced by programs" + + (f" (+{len(rule_dependencies)} dependencies)" if rule_dependencies else "")) # fowl predicates are class-specific (fowl_) and derived from a shared # SMARTS table. Gather the patterns for the classes being predicted so the same @@ -347,7 +306,8 @@ def build_ilp_preds_tensor( worker_state = dict( programs_str=programs_str, col_of=col_of, predicate_set=predicate_set, aux_timeout=aux_timeout, label_timeout=label_timeout, fowl_smarts=fowl_smarts, - rule_programs=rule_programs, computed_facts=computed_facts, + rule_programs=rule_programs, rule_dependencies=rule_dependencies, + computed_facts=computed_facts, ) # Ship mols as RDKit binary with all properties so nothing (stereo/CIP perception, @@ -494,16 +454,15 @@ def predict_smiles( seen.setdefault(pred.name, pred) aux_predicates = list(seen.values()) - # llm_generated_rules: qualify the aux rule references in the rule set per class and - # gather the matching rule programs so the background recomputes their extensions. - rule_programs = None + # llm_generated_rules: gather the rule programs the target classes use (and the library + # programs they build on) so the background recomputes their extensions. + rule_programs = rule_dependencies = None if predicate_set == "llm_generated_rules": programs_by_class = { t[len("chebi_"):]: "\n".join(r for r in rules if f"chebi_{t[len('chebi_'):]}" in r) for t in target_predicates if t.startswith("chebi_") } - rules_qualified, rule_programs = _qualify_aux_rules(programs_by_class, aux_library_dir) - rules = "\n".join(rules_qualified.values()).split("\n") + rule_programs, rule_dependencies = _collect_rule_programs(programs_by_class, aux_library_dir) # fowl predicates are class-specific: gather the SMARTS patterns for the target # classes so the background emits the fowl_ facts the rules reference. @@ -527,7 +486,8 @@ def predict_smiles( background_facts = build_full_background( mol_df, predicate_set=predicate_set, aux_predicates=aux_predicates, aux_timeout=aux_timeout, - fowl_smarts=fowl_smarts, rule_programs=rule_programs, computed_facts=computed_facts, + fowl_smarts=fowl_smarts, rule_programs=rule_programs, + rule_dependencies=rule_dependencies, computed_facts=computed_facts, ) print(f"Evaluating {smiles!r} against {len(rules)} rules and {len(background_facts)} background facts...") print(f" Target predicates: {', '.join(target_predicates)}") diff --git a/chebILP/ilp_problem_builder.py b/chebILP/ilp_problem_builder.py index e89fe4c..407f246 100644 --- a/chebILP/ilp_problem_builder.py +++ b/chebILP/ilp_problem_builder.py @@ -1,506 +1,586 @@ -import os -from typing import Literal -import networkx as nx - -import tqdm -from chebILP.molecule_processing.data_preparation import ChEBIDataset -from chebILP.molecule_processing.mol_to_fol import mol_to_fol_fgs -from chebi_utils.extract_properties import mol_to_fol_atoms, get_numerical_facts -from chebILP.predicate_generation.auxiliary_predicates import load_auxiliary_predicates, compute_auxiliary_extensions, DEFAULT_AUX_TIMEOUT -from chebILP.predicate_generation.auxiliary_rules import derive_rule_extensions, load_class_rules -from chebILP.molecule_processing.fg_matching import get_chembl_fgs, get_chebi_fgs -from chebILP.molecule_processing.fowl_predicates import build_fowl_predicate, calculate_fowl_predicate -import pandas as pd -from chebILP.utils import AVAILABLE_PREDICATE_SETS, get_atom_id -from chebILP.ilp_path_manager import get_bk_path, get_bias_path, get_exs_path -from chebILP.evaluation.clingo_eval import evaluate_with_clingo -from chebi_utils.sample_filters import get_direct_neighbors - - -CHEBI_FG_RULES_PATH = os.path.join("data", "chebi_fg_rules_from_smiles.pl") -CHEBI_FG_LEARNED_RULES_PATH = os.path.join("data", "chebi_fg_learned_rules.pl") -# SMARTS patterns (one per ChEBI class that has a wildcard-bearing molecule) used -# by the "fowl" predicate set, produced by fowl_predicates.build_fowl_smarts. -FOWL_SMARTS_PATH = os.path.join("data", "fowl_smarts.csv") - - -def load_fowl_smarts(path=FOWL_SMARTS_PATH) -> dict[str, str]: - """Load the fowl SMARTS CSV (``chebi_id,SMARTS``) into ``{chebi_id: smarts}``. - - Returns an empty dict if the file is missing so ``build_bk`` degrades to the - plain atom predicates for classes without a fowl pattern. Only a subset of - classes have an entry (those whose molecule carries a ``*``/R wildcard). - """ - if not os.path.exists(path): - return {} - mapping: dict[str, str] = {} - with open(path, "r") as f: - next(f, None) # skip header - for line in f: - line = line.rstrip("\n") - if not line.strip(): - continue - chebi_id, smarts = line.split(",", 1) - mapping[chebi_id.strip()] = smarts.strip() - return mapping - -class ILPProblemBuilder: - - def __init__(self, chebi_version: int, three_star_only: bool = True, base_dir: str = "data", min_pos_samples: int = 25, predicate_set: AVAILABLE_PREDICATE_SETS = "atoms", aux_timeout: float = DEFAULT_AUX_TIMEOUT, aux_library_dir: str | None = None, computed_facts: bool = True): - self.predicate_set = predicate_set - self.problem_dir = os.path.join(base_dir, "ilp_problems") - os.makedirs(self.problem_dir, exist_ok=True) - # Per-call wall-clock budget for LLM-generated auxiliary predicates. - self.aux_timeout = aux_timeout - self.aux_library_dir = aux_library_dir - # When set, molecular-weight and ring-size facts are computed and made - # available to llm_generated_rules during rule evaluation, but never - # written to bk.pl (only the derived aux_* extensions are). - self.computed_facts = computed_facts - - # --- Load pre-built ChEBI data ------------------------------------- - self.dataset = ChEBIDataset(chebi_version=chebi_version, three_star_only=three_star_only, base_dir=base_dir, min_pos_samples=min_pos_samples) - self.hierarchy_graph = nx.transitive_closure_dag(self.dataset.chebi_graph) - self.splits = self.dataset.load_splits_from_csv() - - - def build_examples(self, target_ids: list[str], min_pos_samples=25, max_pos_samples=200, min_neg_samples=25, max_neg_samples=200): - min_n_pos = max_pos_samples + 1 - min_n_pos_id = None - min_n_neg = max_neg_samples + 1 - min_n_neg_id = None - for target_id in tqdm.tqdm(target_ids, desc="Building examples for ChEBI classes"): - n_pos, n_neg = self.gather_samples_for_chebi_cls(target_id, min_pos_samples, max_pos_samples, min_neg_samples, max_neg_samples) - if n_pos < min_n_pos: - min_n_pos = n_pos - min_n_pos_id = target_id - if n_neg < min_n_neg: - min_n_neg = n_neg - min_n_neg_id = target_id - print(f"Label with least positive samples: ChEBI:{min_n_pos_id} with {min_n_pos} samples") - print(f"Label with least negative samples: ChEBI:{min_n_neg_id} with {min_n_neg} samples") - - - def build_bk(self, target_ids): - """ - Build ILP background knowledge. - - Args: - """ - - rules, rule_predicates = [], [] - if self.predicate_set in ["chebi_fg_rules", "chebi_fg_learned_rules"]: - prolog_lines_rules, body_predicates_rules = build_background_chebi_fg_rules(CHEBI_FG_RULES_PATH if self.predicate_set == "chebi_fg_rules" else CHEBI_FG_LEARNED_RULES_PATH) - rules = prolog_lines_rules - rule_predicates = body_predicates_rules - - for target_id in tqdm.tqdm(target_ids, desc="Building background knowledge for ChEBI classes"): - print(f"Building background knowledge for ChEBI:{target_id}...") - - # LLM-generated auxiliary predicates are specific to the target class, - # so they are loaded once per target and merged into the atom-level BK. - aux_predicates = None - if self.predicate_set == "llm_generated_fgs": - aux_predicates = load_auxiliary_predicates(target_id, library_dir=self.aux_library_dir) - print(f" Loaded {len(aux_predicates)} auxiliary predicate(s) for ChEBI:{target_id}") - - # llm_generated_rules: the class's auxiliary predicates are ASP rules, - # evaluated (below) against the atom facts plus optional computed facts. - # Only the derived aux_* extensions are written to bk.pl. - rule_programs = None - if self.predicate_set == "llm_generated_rules": - rule_programs = load_class_rules(target_id, library_dir=self.aux_library_dir) - print(f" Loaded {len(rule_programs)} auxiliary rule(s) for ChEBI:{target_id}") - - # The fowl set adds a single class-specific predicate, fowl_, - # derived from a SMARTS pattern, on top of the atom predicates. Not every - # class has a pattern; those fall back to the plain atom predicates. - fowl_smarts = None - if self.predicate_set == "fowl": - if not hasattr(self, "_fowl_smarts"): - self._fowl_smarts = load_fowl_smarts() - smarts = self._fowl_smarts.get(target_id) - if smarts is None: - print(f" No fowl SMARTS for ChEBI:{target_id}; falling back to plain atom predicates.") - else: - print(f" Loaded fowl SMARTS for ChEBI:{target_id}: {smarts}") - fowl_smarts = {target_id: smarts} - - selected_ids_by_split = dict() - prolog_lines_by_split = dict() - computed_lines_by_split = dict() - body_predicates = set() - for split in ["train", "validation", "test"]: - exs_path = get_exs_path(target_id, base_dir=self.problem_dir, split=split) - with open(exs_path, "r") as f: - # for each line get id between inner parentheses (e.g. pos(chebi_123(456)). -> 456) and select corresponding rows from samples_df - selected_ids = [line.strip().split("(")[-1].split(")")[0] for line in f.readlines() if line.strip() and not line.startswith("%")] - selected_rows = self.dataset.molecules[[id in selected_ids for id in self.dataset.molecules.index]] - selected_ids_by_split[split] = selected_ids - - # standard bk is always added - prolog_lines = [] - prolog_lines_atoms, body_predicates_atoms = build_background_chemlog(selected_rows, aux_predicates=aux_predicates, aux_timeout=self.aux_timeout, predicate_set=self.predicate_set, fowl_smarts=fowl_smarts) - prolog_lines += prolog_lines_atoms - body_predicates.update(body_predicates_atoms) - if self.predicate_set in ["chembl_fgs", "chebi_fgs"]: - # add fgs as samples - if not hasattr(self, "_fg_data"): - if self.predicate_set == "chembl_fgs": - self._fg_data = get_chembl_fgs(self.dataset.molecules) - else: - self._fg_data = get_chebi_fgs(self.dataset.molecules) - prolog_lines_fgs, body_predicates_fgs = build_background_fg_data(self._fg_data, selected_rows, source=self.predicate_set) - prolog_lines += prolog_lines_fgs - body_predicates.update(body_predicates_fgs) - prolog_lines_by_split[split] = prolog_lines - - # Computed facts (molecular weight, ring size) feed rule evaluation - # only; they are intentionally kept out of prolog_lines (bk.pl). - if self.predicate_set == "llm_generated_rules" and self.computed_facts: - computed_lines_by_split[split] = build_computed_facts(selected_rows) - - # for evaluating rules, merge alls splits, separate results afterwards - if self.predicate_set in ["chebi_fg_rules", "chebi_fg_learned_rules"]: - all_selected_ids = [id for split in ["train", "validation", "test"] for id in selected_ids_by_split[split]] - all_prolog_lines = [line for split in ["train", "validation", "test"] for line in prolog_lines_by_split[split]] - positives = evaluate_with_clingo(rules, all_prolog_lines, rule_predicates, all_selected_ids, list(body_predicates)) - for positive_extension in positives: - pred = positive_extension - in_split = {"train": False, "validation": False, "test": False} - for example in positives[positive_extension]: - for split in ["train", "validation", "test"]: - if example in selected_ids_by_split[split]: - if not in_split[split]: - body_predicates.add((pred, 1)) - in_split[split] = True - prolog_lines_by_split[split].append(f"{pred}({example}).") - - # llm_generated_rules: ground each class rule against the atom facts plus - # computed facts (all splits merged), then write only the derived aux_* - # extensions back into each split. The computed facts are never written. - if self.predicate_set == "llm_generated_rules" and rule_programs: - all_selected_ids = [id for split in ["train", "validation", "test"] for id in selected_ids_by_split[split]] - eval_facts = [line for split in ["train", "validation", "test"] for line in prolog_lines_by_split[split]] - if self.computed_facts: - eval_facts += [line for split in ["train", "validation", "test"] for line in computed_lines_by_split.get(split, [])] - # The class's rules are grounded as one program, so a rule may use a predicate - # another of its rules defines. The head may be of any arity; each derived - # atom is written to the split of the molecule it belongs to. - extensions = derive_rule_extensions(rule_programs, eval_facts, all_selected_ids) - for rp in rule_programs: - emitted = {split: set() for split in ["train", "validation", "test"]} - for example, arg_tuples in extensions.get(rp.name, {}).items(): - for split in ["train", "validation", "test"]: - if example not in selected_ids_by_split[split]: - continue - for args in arg_tuples: - line = f"{rp.name}({','.join(args)})." - if line in emitted[split]: - continue - emitted[split].add(line) - body_predicates.add((rp.name, len(args))) - prolog_lines_by_split[split].append(line) - - for split in ["train", "validation", "test"]: - prolog_lines = prolog_lines_by_split[split] - bk_path = get_bk_path(target_id, base_dir=self.problem_dir, predicate_set=self.predicate_set, split=split) - - with open(bk_path, "w+") as f: - f.write("\n".join(prolog_lines) + "\n") - - # create bias file template based on bk predicates - plain_bias_path = get_bias_path(target_id, split="train", base_dir=self.problem_dir, predicate_set=self.predicate_set) # bias file path for settings-specific bias file (created in build_bias) - bias_lines = [ - f"%% CHEBI:{target_id} (bias file without settings)", - f"", - f"%% max_vars(TODO).", - f"%% max_body(TODO).", - f"%% max_clauses(TODO).", - f"", - f"head_pred(chebi_{target_id}, 1)."] + [ - f"body_pred({pred},{arity})." for pred, arity in body_predicates - ] - # bias without settings (as template) - with open(plain_bias_path, "w+") as f: - f.write("\n".join(bias_lines) + "\n") - - - def build_negative_mix(self, neg_pool: pd.DataFrame, sibling_ids: set, max_samples: int, random_state: int = 42) -> pd.DataFrame: - """50:50 mix of direct-sibling negatives and random negatives from ``neg_pool``. - - Up to half of ``max_samples`` are the target's direct siblings (near-misses); the rest - are drawn uniformly at random from the non-sibling remainder. When a class has fewer - siblings than half, the random draw takes up the slack rather than the set shrinking, so - the objective is global classification instead of separation from the superclass alone. - """ - half = max_samples // 2 - sibling_negs = neg_pool[neg_pool.index.astype(str).isin(sibling_ids)] - if len(sibling_negs) > half: - sibling_negs = sibling_negs.sample(half, random_state=random_state) - random_pool = neg_pool[~neg_pool.index.astype(str).isin(sibling_ids)] - n_random = min(max_samples - len(sibling_negs), len(random_pool)) - random_negs = random_pool.sample(n_random, random_state=random_state) if n_random > 0 else random_pool.iloc[:0] - return pd.concat([sibling_negs, random_negs]) - - def gather_samples_for_chebi_cls(self, target_id: str, min_pos_samples=25, max_pos_samples=200, min_neg_samples=25, max_neg_samples=200): - descendants = list(self.hierarchy_graph.predecessors(target_id)) + [target_id] - # not all descendants are molecules (i.e., have a SMILES annotation) -> only take the ones that are in the samples_df (i.e. have a SMILES annotation and are in the 3_STAR subset) - - df_pos = self.dataset.molecules[[id in descendants for id in self.dataset.molecules.index]] - df_neg = self.dataset.molecules[[id not in df_pos.index for id in self.dataset.molecules.index]] - if len(df_pos) < min_pos_samples: - print(f"ChEBI class {target_id} does not have enough positive samples (found {len(df_pos)}, required are at least {min_pos_samples}). Got samples {df_pos.index.tolist()}") - if len(df_neg) < min_neg_samples: - print(f"ChEBI class {target_id} does not have enough negative samples (found {len(df_neg)}, required are at least {min_neg_samples}). Got samples {df_neg.index.tolist()}") - - # Direct-sibling molecules: subclasses shared with the target's parents. They form the - # near-miss half of every split's negatives; the other half is drawn uniformly at random - # from the full negative pool. The objective is therefore global classification, not - # separating the target from its superclass only. - mol_index = set(str(i) for i in self.dataset.molecules.index) - pos_ids, sibling_neg_ids = get_direct_neighbors(mol_index, self.dataset.chebi_graph, target_id) - sibling_neg_ids = set(sibling_neg_ids) - - samples_by_split = dict() - pos_train_samples = df_pos[df_pos.index.astype(str).isin(self.splits[self.splits["split"] == "train"])] - samples_by_split[("pos", "train")] = pos_train_samples.sample(min(max_pos_samples, len(pos_train_samples)), random_state=42) # if there are more positives than max_pos_samples, sample randomly - neg_train_samples = df_neg[df_neg.index.astype(str).isin(self.splits[self.splits["split"] == "train"])] - samples_by_split[("neg", "train")] = self.build_negative_mix(neg_train_samples, sibling_neg_ids, max_neg_samples) - - samples_by_split[("pos", "validation")] = df_pos[df_pos.index.astype(str).isin(self.splits[self.splits["split"] == "validation"]) & df_pos.index.astype(str).isin(pos_ids)] - neg_val_samples = df_neg[df_neg.index.astype(str).isin(self.splits[self.splits["split"] == "validation"])] - samples_by_split[("neg", "validation")] = self.build_negative_mix(neg_val_samples, sibling_neg_ids, max_neg_samples) - samples_by_split[("pos", "test")] = df_pos[df_pos.index.astype(str).isin(self.splits[self.splits["split"] == "test"]) & df_pos.index.astype(str).isin(pos_ids)] - neg_test_samples = df_neg[df_neg.index.astype(str).isin(self.splits[self.splits["split"] == "test"])] - samples_by_split[("neg", "test")] = self.build_negative_mix(neg_test_samples, sibling_neg_ids, max_neg_samples) - - for (posneg, split), df in samples_by_split.items(): - exs_path = get_exs_path(target_id, base_dir=self.problem_dir, split=split) - with open(exs_path, "w+" if posneg == "pos" else "a") as f: - for sample in df.index: - f.write(f"{posneg}(chebi_{target_id}({sample})).\n") - - # sum up all positive and negative samples across splits - return sum(len(v) for k, v in samples_by_split.items() if k[0] == "pos"), sum(len(v) for k, v in samples_by_split.items() if k[0] == "neg") - - - -def build_background_chemlog(rows, aux_predicates=None, aux_timeout=DEFAULT_AUX_TIMEOUT, aux_failures=None, predicate_set="atoms", fowl_smarts=None): - comments = [] - lines_by_predicate, arities = {}, {} - if "farm_fgs" in predicate_set: - lines_by_predicate["has_fg"] = [] - arities["has_fg"] = 2 - if "atoms" in predicate_set or "farm_fgs" not in predicate_set: - lines_by_predicate["has_atom"] = [] - arities["has_atom"] = 2 - - aux_ext_by_mol = {} - if aux_predicates: - aux_ext_by_mol = compute_auxiliary_extensions( - aux_predicates, - [(row.Index, row.mol) for row in rows.itertuples()], - timeout=aux_timeout, - failures=aux_failures, - ) - - for row in rows.itertuples(): - atom_extensions, fg_extensions, mol_extensions = {}, {}, set() - if "farm_fgs" in predicate_set: - # Functional-group level model: entities are FARM functional-group - # nodes rather than atoms. has_fg links the molecule to its FG nodes. - fg_extensions = mol_to_fol_fgs(row.mol, add_fg_atom_predicates="atoms" in predicate_set) - node_ids = sorted({id for ids in fg_extensions.values() for nid in ids for id in (nid if isinstance(nid, tuple) else (nid,))}) # flatten tuples - for node_id in node_ids: - if node_id >= row.mol.GetNumAtoms(): - fg_id = get_atom_id(node_id, row.Index) - lines_by_predicate["has_fg"].append( - f"has_fg({row.Index},{fg_id}).") - if "atoms" in predicate_set or "farm_fgs" not in predicate_set: - for atom in row.mol.GetAtoms(): - atom_id = get_atom_id(atom.GetIdx(), row.Index) - lines_by_predicate["has_atom"].append(f"has_atom({row.Index},{atom_id}).") - - atom_extensions, mol_extensions = mol_to_fol_atoms(row.mol) - - # Merge LLM-generated auxiliary predicates. Their names are ``aux_``-prefixed, - # so they never collide with the built-in extensions produced above. - if aux_predicates: - aux_atom_ext, aux_mol_ext = aux_ext_by_mol.get(row.Index, ({}, set())) - atom_extensions.update(aux_atom_ext) - mol_extensions.update(aux_mol_ext) - - # fowl: class-specific SMARTS-match predicates (fowl_) added on - # top of the atom predicates. Each match binds the pattern's wildcard - # atoms, so the arity equals the number of wildcards; the tuples are - # emitted as atom-id arguments by the extension loop below. - if fowl_smarts: - for cls_id, smarts in fowl_smarts.items(): - predicate_name, _ = build_fowl_predicate(smarts, cls_id) - try: - matches = calculate_fowl_predicate(smarts, row.mol) - except Exception as e: - print(f"Warning: failed to compute {predicate_name} for CHEBI:{row.Index}: {e}") - continue - if matches: - atom_extensions.setdefault(predicate_name, []).extend(matches) - - for predicate, indices in {**atom_extensions, **fg_extensions}.items(): - if predicate.startswith("cip_code_"): - predicate = "cip_code_" + predicate[-1].upper() - if (predicate in {"EQ", "atom", "*", "r", "r#"} or (predicate.startswith("r") and predicate[1:].isdigit() and int(predicate[1:]) > 0) or not indices): - continue - - is_tuple = isinstance(indices[0], tuple) - if predicate not in lines_by_predicate: - lines_by_predicate[predicate] = [] - if predicate not in arities: - arities[predicate] = len(indices[0]) if is_tuple else 1 - if is_tuple: - for args in indices: - arg_str = ",".join(get_atom_id(a, row.Index) for a in args) - lines_by_predicate[predicate].append(f"{predicate}({arg_str}).") - else: - for idx in indices: - lines_by_predicate[predicate].append(f"{predicate}({get_atom_id(idx, row.Index)}).") - - for predicate in mol_extensions: - if predicate not in lines_by_predicate: - lines_by_predicate[predicate] = [] - if predicate not in arities: - arities[predicate] = 1 - lines_by_predicate[predicate].append(f"{predicate}({row.Index}).") - - return comments + [line for lines in lines_by_predicate.values() for line in lines], [(pred, arities[pred]) for pred in arities.keys()] - - -def build_computed_facts(rows): - """Molecular-weight and ring-size facts used only to evaluate llm_generated_rules. - - Formats ``chebi_utils.get_numerical_facts`` per molecule as Prolog facts - (``mol_weight(Mol, W)``, one ``ring_size(Mol, Size)`` per ring). These facts are - fed to Clingo when a class's auxiliary rules are grounded, but are never written to - ``bk.pl`` — only the derived ``aux_*`` extensions are persisted. - """ - lines = [] - for row in rows.itertuples(): - for pred, values in get_numerical_facts(row.mol).items(): - for value in values: - lines.append(f"{pred}({row.Index},{value}).") - return lines - - -def build_full_background( - rows: pd.DataFrame, - predicate_set: AVAILABLE_PREDICATE_SETS = "atoms", - aux_predicates=None, - aux_timeout: float = DEFAULT_AUX_TIMEOUT, - aux_failures=None, - fowl_smarts=None, - rule_programs=None, - computed_facts: bool = True, -) -> list[str]: - """Build one flat background-knowledge fact list for the molecules in ``rows``. - - Mirrors :meth:`ILPProblemBuilder.build_bk` so prediction tensors are evaluated - against exactly the same BK the programs were learned on, rather than always the - plain ``atoms`` set. For the ``chebi_fg_rules`` / ``chebi_fg_learned_rules`` sets the - functional-group rule clauses are added to the BK directly (rather than pre-evaluated - into facts): the caller grounds and solves the combined program, which derives them. - - All work is scoped to ``rows``, so this can be called per molecule (e.g. to bound - Clingo grounding memory). ``aux_predicates`` (for ``llm_generated_fgs``) are the - name-deduplicated predicates gathered across all classes; their extensions are - evaluated on ``rows`` here. - """ - prolog_lines, _ = build_background_chemlog( - rows, aux_predicates=aux_predicates, aux_timeout=aux_timeout, aux_failures=aux_failures, - predicate_set=predicate_set, fowl_smarts=fowl_smarts, - ) - prolog_lines = list(prolog_lines) - - if predicate_set in ("chembl_fgs", "chebi_fgs"): - fg_data = get_chembl_fgs(rows) if predicate_set == "chembl_fgs" else get_chebi_fgs(rows) - fg_lines, _ = build_background_fg_data(fg_data, rows, source=predicate_set) - prolog_lines += fg_lines - - if predicate_set in ("chebi_fg_rules", "chebi_fg_learned_rules"): - rule_lines, _ = build_background_chebi_fg_rules( - CHEBI_FG_RULES_PATH if predicate_set == "chebi_fg_rules" else CHEBI_FG_LEARNED_RULES_PATH - ) - prolog_lines += rule_lines - - # llm_generated_rules: recompute the class's aux_* extensions exactly as build_bk - # does (ground each rule over atom + computed facts) and append them as facts, so a - # learned program's aux_* body literals resolve. Computed facts stay local to the - # grounding and are not added to the returned BK. - if predicate_set == "llm_generated_rules" and rule_programs: - eval_facts = list(prolog_lines) - if computed_facts: - eval_facts += build_computed_facts(rows) - mol_ids = [str(i) for i in rows.index] - extensions = derive_rule_extensions(rule_programs, eval_facts, mol_ids) - for rp in rule_programs: - emitted = set() - for arg_tuples in extensions.get(rp.name, {}).values(): - for args in arg_tuples: - line = f"{rp.name}({','.join(args)})." - if line not in emitted: - emitted.add(line) - prolog_lines.append(line) - - return prolog_lines - - -def build_background_chebi_fg_rules(rules_path=None): - """Load ChEBI functional group rules from a Prolog file and return them as BK lines and body predicates. - - Each rule defines a chebi_XXXXX(M) predicate in terms of atom-level predicates. - These are added as Prolog rules to the BK and as body_pred entries (arity 1) in the bias. - """ - if rules_path is None: - rules_path = CHEBI_FG_RULES_PATH - - prolog_lines = [f"% ChEBI FG rules from {os.path.basename(rules_path)}"] - body_predicates = [] - seen_predicates = set() - - with open(rules_path, "r") as f: - for line in f: - line = line.strip() - if not line or line.startswith("%"): - continue - prolog_lines.append(line) - # Extract predicate name from head: chebi_XXXXX(M) :- ... - pred_name = line.split("(")[0].strip() - if pred_name and pred_name not in seen_predicates: - seen_predicates.add(pred_name) - body_predicates.append(pred_name) - - print(f"Loaded {len(body_predicates)} ChEBI FG rule predicates from {rules_path}") - return prolog_lines, body_predicates - - -def build_background_fg_data(fg_data: dict[int, list[str]], rows, source: Literal["chembl_fgs", "chebi_fgs"]): - lines_by_predicate = dict() - - for row in rows.itertuples(): - if row.Index not in fg_data: - print(f"Warning: No functional group data found for CHEBI:{row.Index} in source {source}. This molecule will only have atom and bond predicates in the background knowledge.") - continue - for fg in fg_data[row.Index]: - if fg not in lines_by_predicate: - lines_by_predicate[fg] = [] - lines_by_predicate[fg].append(f"{fg}({row.Index}).") - total_lines = [line for lines in lines_by_predicate.values() for line in lines] - return total_lines, [(pred, 1) for pred in lines_by_predicate.keys()] - - -if __name__ == "__main__": - builder = ILPProblemBuilder( - chebi_version=251, - predicate_set="atoms", - ) - target_ids = ["134362"] - builder.build_examples(target_ids) +import os +from typing import Literal +import networkx as nx + +import tqdm +from chebILP.molecule_processing.data_preparation import ChEBIDataset +from chebILP.molecule_processing.mol_to_fol import mol_to_fol_fgs +from chebi_utils.extract_properties import mol_to_fol_atoms, get_numerical_facts +from chebILP.predicate_generation.auxiliary_predicates import load_auxiliary_predicates, compute_auxiliary_extensions, DEFAULT_AUX_TIMEOUT +from chebILP.predicate_generation.auxiliary_rules import derive_rule_extensions, load_class_rules, resolve_rule_dependencies +from chebILP.molecule_processing.fg_matching import get_chembl_fgs, get_chebi_fgs +from chebILP.molecule_processing.fowl_predicates import build_fowl_predicate, calculate_fowl_predicate +import pandas as pd +from chebILP.utils import AVAILABLE_PREDICATE_SETS, get_atom_id +from chebILP.ilp_path_manager import get_bk_path, get_bias_path, get_exs_path +from chebILP.evaluation.clingo_eval import evaluate_with_clingo + + +CHEBI_FG_RULES_PATH = os.path.join("data", "chebi_fg_rules_from_smiles.pl") +CHEBI_FG_LEARNED_RULES_PATH = os.path.join("data", "chebi_fg_learned_rules.pl") +# SMARTS patterns (one per ChEBI class that has a wildcard-bearing molecule) used +# by the "fowl" predicate set, produced by fowl_predicates.build_fowl_smarts. +FOWL_SMARTS_PATH = os.path.join("data", "fowl_smarts.csv") + + +def load_fowl_smarts(path=FOWL_SMARTS_PATH) -> dict[str, str]: + """Load the fowl SMARTS CSV (``chebi_id,SMARTS``) into ``{chebi_id: smarts}``. + + Returns an empty dict if the file is missing so ``build_bk`` degrades to the + plain atom predicates for classes without a fowl pattern. Only a subset of + classes have an entry (those whose molecule carries a ``*``/R wildcard). + """ + if not os.path.exists(path): + return {} + mapping: dict[str, str] = {} + with open(path, "r") as f: + next(f, None) # skip header + for line in f: + line = line.rstrip("\n") + if not line.strip(): + continue + chebi_id, smarts = line.split(",", 1) + mapping[chebi_id.strip()] = smarts.strip() + return mapping + +class ILPProblemBuilder: + + def __init__(self, chebi_version: int, three_star_only: bool = True, base_dir: str = "data", min_pos_samples: int = 25, predicate_set: AVAILABLE_PREDICATE_SETS = "atoms", aux_timeout: float = DEFAULT_AUX_TIMEOUT, aux_library_dir: str | None = None, computed_facts: bool = True): + self.predicate_set = predicate_set + self.problem_dir = os.path.join(base_dir, "ilp_problems") + os.makedirs(self.problem_dir, exist_ok=True) + # Per-call wall-clock budget for LLM-generated auxiliary predicates. + self.aux_timeout = aux_timeout + self.aux_library_dir = aux_library_dir + # When set, molecular-weight and ring-size facts are computed and made + # available to llm_generated_rules during rule evaluation, but never + # written to bk.pl (only the derived aux_* extensions are). + self.computed_facts = computed_facts + + # --- Load pre-built ChEBI data ------------------------------------- + self.dataset = ChEBIDataset(chebi_version=chebi_version, three_star_only=three_star_only, base_dir=base_dir, min_pos_samples=min_pos_samples) + self.hierarchy_graph = nx.transitive_closure_dag(self.dataset.chebi_graph) + self.splits = self.dataset.load_splits_from_csv() + + # Invariant across target classes, so built once rather than per class. + self._mol_index = set(self.dataset.molecules.index) + self._split_ids = {split: set(self.splits[self.splits["split"] == split]["id"].astype(str)) + for split in ["train", "validation", "test"]} + # Size of the molecule graph, used to prefer small molecules when a split is + # capped. This is the same count that drives the has_atom facts in bk.pl, so it + # includes explicit hydrogens where a molecule carries them. + self._atom_counts = self.dataset.molecules["mol"].map(lambda m: m.GetNumAtoms()) + + + def build_examples(self, target_ids: list[str], min_pos_samples=25, max_pos_samples=200, min_neg_samples=25, max_neg_samples=200): + # Counts are summed over the three splits, so they are not comparable against the + # per-split max_*_samples caps; take the minimum over what was actually written. + counts = {} + for target_id in tqdm.tqdm(target_ids, desc="Building examples for ChEBI classes"): + counts[target_id] = self.gather_samples_for_chebi_cls(target_id, min_pos_samples, max_pos_samples, min_neg_samples, max_neg_samples) + if not counts: + return + min_n_pos_id = min(counts, key=lambda c: counts[c][0]) + min_n_neg_id = min(counts, key=lambda c: counts[c][1]) + print(f"Label with least positive samples: ChEBI:{min_n_pos_id} with {counts[min_n_pos_id][0]} samples across all splits") + print(f"Label with least negative samples: ChEBI:{min_n_neg_id} with {counts[min_n_neg_id][1]} samples across all splits") + + + def build_bk(self, target_ids): + """ + Build ILP background knowledge. + + Args: + """ + + rules, rule_predicates = [], [] + failed_rule_classes: list[str] = [] + if self.predicate_set in ["chebi_fg_rules", "chebi_fg_learned_rules"]: + prolog_lines_rules, body_predicates_rules = build_background_chebi_fg_rules(CHEBI_FG_RULES_PATH if self.predicate_set == "chebi_fg_rules" else CHEBI_FG_LEARNED_RULES_PATH) + rules = prolog_lines_rules + rule_predicates = body_predicates_rules + + pbar = tqdm.tqdm(target_ids, desc="Building background knowledge") + for target_id in pbar: + # The per-class status goes into the bar itself; printing it would redraw the bar + # on every iteration. Only warnings and failures are written as their own lines. + pbar.set_description(f"Building background knowledge for ChEBI:{target_id}") + pbar.set_postfix_str("") + + # LLM-generated auxiliary predicates are specific to the target class, + # so they are loaded once per target and merged into the atom-level BK. + aux_predicates = None + if self.predicate_set == "llm_generated_fgs": + aux_predicates = load_auxiliary_predicates(target_id, library_dir=self.aux_library_dir) + pbar.set_postfix_str(f"{len(aux_predicates)} aux predicate(s)") + + # llm_generated_rules: the class's auxiliary predicates are ASP rules, + # evaluated (below) against the atom facts plus optional computed facts. + # Only the derived aux_* extensions are written to bk.pl. + rule_programs, dependency_programs = None, [] + if self.predicate_set == "llm_generated_rules": + rule_programs = load_class_rules(target_id, library_dir=self.aux_library_dir) + # class_map.json records only the predicates the class chose, not the ones + # they build on, so the dependencies have to be pulled in from the library + # or the rules ground against an empty body and derive nothing. + dependency_programs = resolve_rule_dependencies(rule_programs, self.aux_library_dir) + pbar.set_postfix_str(f"{len(rule_programs)} rule(s)" + + (f" +{len(dependency_programs)} dep(s)" if dependency_programs else "")) + + # The fowl set adds a single class-specific predicate, fowl_, + # derived from a SMARTS pattern, on top of the atom predicates. Not every + # class has a pattern; those fall back to the plain atom predicates. + fowl_smarts = None + if self.predicate_set == "fowl": + if not hasattr(self, "_fowl_smarts"): + self._fowl_smarts = load_fowl_smarts() + smarts = self._fowl_smarts.get(target_id) + if smarts is None: + pbar.set_postfix_str("no fowl SMARTS, plain atom predicates") + else: + pbar.set_postfix_str(f"fowl SMARTS {smarts}") + fowl_smarts = {target_id: smarts} + + selected_ids_by_split = dict() + prolog_lines_by_split = dict() + computed_lines_by_split = dict() + body_predicates = set() + for split in ["train", "validation", "test"]: + exs_path = get_exs_path(target_id, base_dir=self.problem_dir, split=split) + with open(exs_path, "r") as f: + # for each line get id between inner parentheses (e.g. pos(chebi_123(456)). -> 456) and select corresponding rows from samples_df + selected_ids = [line.strip().split("(")[-1].split(")")[0] for line in f.readlines() if line.strip() and not line.startswith("%")] + selected_rows = self.dataset.molecules[[id in selected_ids for id in self.dataset.molecules.index]] + selected_ids_by_split[split] = selected_ids + + # standard bk is always added + prolog_lines = [] + prolog_lines_atoms, body_predicates_atoms = build_background_chemlog(selected_rows, aux_predicates=aux_predicates, aux_timeout=self.aux_timeout, predicate_set=self.predicate_set, fowl_smarts=fowl_smarts) + prolog_lines += prolog_lines_atoms + body_predicates.update(body_predicates_atoms) + if self.predicate_set in ["chembl_fgs", "chebi_fgs"]: + # add fgs as samples + if not hasattr(self, "_fg_data"): + if self.predicate_set == "chembl_fgs": + self._fg_data = get_chembl_fgs(self.dataset.molecules) + else: + self._fg_data = get_chebi_fgs(self.dataset.molecules) + prolog_lines_fgs, body_predicates_fgs = build_background_fg_data(self._fg_data, selected_rows, source=self.predicate_set) + prolog_lines += prolog_lines_fgs + body_predicates.update(body_predicates_fgs) + prolog_lines_by_split[split] = prolog_lines + + # Computed facts (molecular weight, ring size) feed rule evaluation + # only; they are intentionally kept out of prolog_lines (bk.pl). + if self.predicate_set == "llm_generated_rules" and self.computed_facts: + computed_lines_by_split[split] = build_computed_facts(selected_rows) + + # for evaluating rules, merge alls splits, separate results afterwards + if self.predicate_set in ["chebi_fg_rules", "chebi_fg_learned_rules"]: + all_selected_ids = [id for split in ["train", "validation", "test"] for id in selected_ids_by_split[split]] + all_prolog_lines = [line for split in ["train", "validation", "test"] for line in prolog_lines_by_split[split]] + positives = evaluate_with_clingo(rules, all_prolog_lines, rule_predicates, all_selected_ids, list(body_predicates)) + for positive_extension in positives: + pred = positive_extension + in_split = {"train": False, "validation": False, "test": False} + for example in positives[positive_extension]: + for split in ["train", "validation", "test"]: + if example in selected_ids_by_split[split]: + if not in_split[split]: + body_predicates.add((pred, 1)) + in_split[split] = True + prolog_lines_by_split[split].append(f"{pred}({example}).") + + # llm_generated_rules: ground each class rule against the atom facts plus + # computed facts (all splits merged), then write only the derived aux_* + # extensions back into each split. The computed facts are never written. + if self.predicate_set == "llm_generated_rules" and rule_programs: + all_selected_ids = [id for split in ["train", "validation", "test"] for id in selected_ids_by_split[split]] + eval_facts = [line for split in ["train", "validation", "test"] for line in prolog_lines_by_split[split]] + if self.computed_facts: + eval_facts += [line for split in ["train", "validation", "test"] for line in computed_lines_by_split.get(split, [])] + # The class's rules are grounded as one program, so a rule may use a predicate + # another of its rules defines. The head may be of any arity; each derived + # atom is written to the split of the molecule it belongs to. Dependencies + # take part in the grounding but never reach bk.pl. + try: + extensions = derive_rule_extensions( + rule_programs + dependency_programs, eval_facts, all_selected_ids + ) + except (RuntimeError, MemoryError) as e: + # One class's rules must not end a run that is hours long. The class keeps + # its atom-level bk.pl and simply goes without its aux_* extensions. + pbar.write(f" Grounding failed for ChEBI:{target_id} ({e}); " + f"continuing without its auxiliary extensions. " + f"Rules: {', '.join(rp.name for rp in rule_programs)}") + failed_rule_classes.append(target_id) + extensions = {} + for rp in rule_programs: + emitted = {split: set() for split in ["train", "validation", "test"]} + for example, arg_tuples in extensions.get(rp.name, {}).items(): + for split in ["train", "validation", "test"]: + if example not in selected_ids_by_split[split]: + continue + for args in arg_tuples: + line = f"{rp.name}({','.join(args)})." + if line in emitted[split]: + continue + emitted[split].add(line) + body_predicates.add((rp.name, len(args))) + prolog_lines_by_split[split].append(line) + + for split in ["train", "validation", "test"]: + prolog_lines = prolog_lines_by_split[split] + bk_path = get_bk_path(target_id, base_dir=self.problem_dir, predicate_set=self.predicate_set, split=split) + + with open(bk_path, "w+") as f: + f.write("\n".join(prolog_lines) + "\n") + + # create bias file template based on bk predicates + plain_bias_path = get_bias_path(target_id, split="train", base_dir=self.problem_dir, predicate_set=self.predicate_set) # bias file path for settings-specific bias file (created in build_bias) + bias_lines = [ + f"%% CHEBI:{target_id} (bias file without settings)", + f"", + f"%% max_vars(TODO).", + f"%% max_body(TODO).", + f"%% max_clauses(TODO).", + f"", + f"head_pred(chebi_{target_id}, 1)."] + [ + f"body_pred({pred},{arity})." for pred, arity in body_predicates + ] + # bias without settings (as template) + with open(plain_bias_path, "w+") as f: + f.write("\n".join(bias_lines) + "\n") + + if failed_rule_classes: + print(f"\n{len(failed_rule_classes)} class(es) built without their auxiliary rule " + f"extensions because grounding failed: {', '.join(failed_rule_classes)}") + + + def _take_smallest(self, df: pd.DataFrame, max_samples: int) -> pd.DataFrame: + """The ``max_samples`` smallest molecules of ``df``, by atom count. + + Ties resolve by the frame's own order, so the pick is deterministic without a seed. + """ + if len(df) <= max_samples: + return df + return df.loc[self._atom_counts[df.index].nsmallest(max_samples, keep="first").index] + + def _direct_neighbors(self, target_id: str) -> tuple[set[str], set[str]]: + """Molecule ids below ``target_id``, and those shared by all its direct parents. + + The second set is the near-miss pool: descendants of every parent that are not + descendants of the target itself. Equivalent to + ``chebi_utils.sample_filters.get_direct_neighbors``, but reuses ``hierarchy_graph`` + instead of rebuilding the transitive closure once per class. + """ + pos_ids = {str(d) for d in self.hierarchy_graph.predecessors(target_id)} & self._mol_index + parent_spaces = [ + {str(d) for d in self.hierarchy_graph.predecessors(parent)} & self._mol_index + for parent in self.dataset.chebi_graph.successors(target_id) + ] + if not parent_spaces: + return pos_ids, set() + return pos_ids, set.intersection(*parent_spaces) - pos_ids + + def build_negative_mix(self, neg_pool: pd.DataFrame, sibling_ids: set, max_samples: int, random_state: int = 42, prefer_smallest: bool = False) -> pd.DataFrame: + """50:50 mix of direct-sibling negatives and random negatives from ``neg_pool``. + + Up to half of ``max_samples`` are the target's direct siblings (near-misses); the rest + are drawn uniformly at random from the non-sibling remainder. When a class has fewer + siblings than half, the random draw takes up the slack rather than the set shrinking, so + the objective is global classification instead of separation from the superclass alone. + + With ``prefer_smallest``, an over-full sibling half keeps the smallest molecules rather + than a random draw, which shrinks the derived bk.pl. Only the training split sets it; + validation and test stay random so their scores remain size-unbiased. + """ + half = max_samples // 2 + sibling_negs = neg_pool[neg_pool.index.astype(str).isin(sibling_ids)] + if len(sibling_negs) > half: + sibling_negs = self._take_smallest(sibling_negs, half) if prefer_smallest else sibling_negs.sample(half, random_state=random_state) + random_pool = neg_pool[~neg_pool.index.astype(str).isin(sibling_ids)] + n_random = min(max_samples - len(sibling_negs), len(random_pool)) + random_negs = random_pool.sample(n_random, random_state=random_state) if n_random > 0 else random_pool.iloc[:0] + return pd.concat([sibling_negs, random_negs]) + + def gather_samples_for_chebi_cls(self, target_id: str, min_pos_samples=25, max_pos_samples=200, min_neg_samples=25, max_neg_samples=200): + descendants = set(self.hierarchy_graph.predecessors(target_id)) | {target_id} + # not all descendants are molecules (i.e., have a SMILES annotation) -> only take the ones that are in the samples_df (i.e. have a SMILES annotation and are in the 3_STAR subset) + + is_pos = self.dataset.molecules.index.isin(descendants) + df_pos = self.dataset.molecules[is_pos] + df_neg = self.dataset.molecules[~is_pos] + if len(df_pos) < min_pos_samples: + print(f"ChEBI class {target_id} does not have enough positive samples (found {len(df_pos)}, required are at least {min_pos_samples}). Got samples {df_pos.index.tolist()}") + if len(df_neg) < min_neg_samples: + print(f"ChEBI class {target_id} does not have enough negative samples (found {len(df_neg)}, required are at least {min_neg_samples}). Got samples {df_neg.index.tolist()}") + + # Direct-sibling molecules: subclasses shared with the target's parents. They form the + # near-miss half of every split's negatives; the other half is drawn uniformly at random + # from the full negative pool. The objective is therefore global classification, not + # separating the target from its superclass only. Only training takes its near-misses + # smallest-first, so validation and test stay size-unbiased. + pos_ids, sibling_neg_ids = self._direct_neighbors(target_id) + + split_ids = self._split_ids + + samples_by_split = dict() + pos_train_samples = df_pos[df_pos.index.astype(str).isin(split_ids["train"])] + # Over the cap, the smallest molecules are kept: they carry the class just as well + # while keeping bk.pl small enough to ground cheaply. + samples_by_split[("pos", "train")] = self._take_smallest(pos_train_samples, max_pos_samples) + neg_train_samples = df_neg[df_neg.index.astype(str).isin(split_ids["train"])] + samples_by_split[("neg", "train")] = self.build_negative_mix(neg_train_samples, sibling_neg_ids, max_neg_samples, prefer_smallest=True) + + samples_by_split[("pos", "validation")] = df_pos[df_pos.index.astype(str).isin(split_ids["validation"]) & df_pos.index.astype(str).isin(pos_ids)] + neg_val_samples = df_neg[df_neg.index.astype(str).isin(split_ids["validation"])] + samples_by_split[("neg", "validation")] = self.build_negative_mix(neg_val_samples, sibling_neg_ids, max_neg_samples) + samples_by_split[("pos", "test")] = df_pos[df_pos.index.astype(str).isin(split_ids["test"]) & df_pos.index.astype(str).isin(pos_ids)] + neg_test_samples = df_neg[df_neg.index.astype(str).isin(split_ids["test"])] + samples_by_split[("neg", "test")] = self.build_negative_mix(neg_test_samples, sibling_neg_ids, max_neg_samples) + + for (posneg, split), df in samples_by_split.items(): + exs_path = get_exs_path(target_id, base_dir=self.problem_dir, split=split) + with open(exs_path, "w+" if posneg == "pos" else "a") as f: + for sample in df.index: + f.write(f"{posneg}(chebi_{target_id}({sample})).\n") + + # sum up all positive and negative samples across splits + return sum(len(v) for k, v in samples_by_split.items() if k[0] == "pos"), sum(len(v) for k, v in samples_by_split.items() if k[0] == "neg") + + + +def build_background_chemlog(rows, aux_predicates=None, aux_timeout=DEFAULT_AUX_TIMEOUT, aux_failures=None, predicate_set="atoms", fowl_smarts=None): + comments = [] + lines_by_predicate, arities = {}, {} + if "farm_fgs" in predicate_set: + lines_by_predicate["has_fg"] = [] + arities["has_fg"] = 2 + if "atoms" in predicate_set or "farm_fgs" not in predicate_set: + lines_by_predicate["has_atom"] = [] + arities["has_atom"] = 2 + + aux_ext_by_mol = {} + if aux_predicates: + aux_ext_by_mol = compute_auxiliary_extensions( + aux_predicates, + [(row.Index, row.mol) for row in rows.itertuples()], + timeout=aux_timeout, + failures=aux_failures, + ) + + for row in rows.itertuples(): + atom_extensions, fg_extensions, mol_extensions = {}, {}, set() + if "farm_fgs" in predicate_set: + # Functional-group level model: entities are FARM functional-group + # nodes rather than atoms. has_fg links the molecule to its FG nodes. + fg_extensions = mol_to_fol_fgs(row.mol, add_fg_atom_predicates="atoms" in predicate_set) + node_ids = sorted({id for ids in fg_extensions.values() for nid in ids for id in (nid if isinstance(nid, tuple) else (nid,))}) # flatten tuples + for node_id in node_ids: + if node_id >= row.mol.GetNumAtoms(): + fg_id = get_atom_id(node_id, row.Index) + lines_by_predicate["has_fg"].append( + f"has_fg({row.Index},{fg_id}).") + if "atoms" in predicate_set or "farm_fgs" not in predicate_set: + for atom in row.mol.GetAtoms(): + atom_id = get_atom_id(atom.GetIdx(), row.Index) + lines_by_predicate["has_atom"].append(f"has_atom({row.Index},{atom_id}).") + + atom_extensions, mol_extensions = mol_to_fol_atoms(row.mol) + + # Merge LLM-generated auxiliary predicates. Their names are ``aux_``-prefixed, + # so they never collide with the built-in extensions produced above. + if aux_predicates: + aux_atom_ext, aux_mol_ext = aux_ext_by_mol.get(row.Index, ({}, set())) + atom_extensions.update(aux_atom_ext) + mol_extensions.update(aux_mol_ext) + + # fowl: class-specific SMARTS-match predicates (fowl_) added on + # top of the atom predicates. Each match binds the pattern's wildcard + # atoms, so the arity equals the number of wildcards; the tuples are + # emitted as atom-id arguments by the extension loop below. + if fowl_smarts: + for cls_id, smarts in fowl_smarts.items(): + predicate_name, _ = build_fowl_predicate(smarts, cls_id) + try: + matches = calculate_fowl_predicate(smarts, row.mol) + except Exception as e: + print(f"Warning: failed to compute {predicate_name} for CHEBI:{row.Index}: {e}") + continue + if matches: + atom_extensions.setdefault(predicate_name, []).extend(matches) + + for predicate, indices in {**atom_extensions, **fg_extensions}.items(): + if predicate.startswith("cip_code_"): + predicate = "cip_code_" + predicate[-1].upper() + if (predicate in {"EQ", "atom", "*", "r", "r#"} or (predicate.startswith("r") and predicate[1:].isdigit() and int(predicate[1:]) > 0) or not indices): + continue + + is_tuple = isinstance(indices[0], tuple) + if predicate not in lines_by_predicate: + lines_by_predicate[predicate] = [] + if predicate not in arities: + arities[predicate] = len(indices[0]) if is_tuple else 1 + if is_tuple: + for args in indices: + arg_str = ",".join(get_atom_id(a, row.Index) for a in args) + lines_by_predicate[predicate].append(f"{predicate}({arg_str}).") + else: + for idx in indices: + lines_by_predicate[predicate].append(f"{predicate}({get_atom_id(idx, row.Index)}).") + + for predicate in mol_extensions: + if predicate not in lines_by_predicate: + lines_by_predicate[predicate] = [] + if predicate not in arities: + arities[predicate] = 1 + lines_by_predicate[predicate].append(f"{predicate}({row.Index}).") + + return comments + [line for lines in lines_by_predicate.values() for line in lines], [(pred, arities[pred]) for pred in arities.keys()] + + +def build_computed_facts(rows): + """Molecular-weight and ring-size facts used only to evaluate llm_generated_rules. + + Formats ``chebi_utils.get_numerical_facts`` per molecule as Prolog facts + (``mol_weight(Mol, W)``, one ``ring_size(Mol, Size)`` per ring). These facts are + fed to Clingo when a class's auxiliary rules are grounded, but are never written to + ``bk.pl`` — only the derived ``aux_*`` extensions are persisted. + """ + lines = [] + for row in rows.itertuples(): + for pred, values in get_numerical_facts(row.mol).items(): + for value in values: + lines.append(f"{pred}({row.Index},{value}).") + return lines + + +def build_full_background( + rows: pd.DataFrame, + predicate_set: AVAILABLE_PREDICATE_SETS = "atoms", + aux_predicates=None, + aux_timeout: float = DEFAULT_AUX_TIMEOUT, + aux_failures=None, + fowl_smarts=None, + rule_programs=None, + rule_dependencies=None, + computed_facts: bool = True, + aux_library_dir: str | None = None, +) -> list[str]: + """Build one flat background-knowledge fact list for the molecules in ``rows``. + + Mirrors :meth:`ILPProblemBuilder.build_bk` so prediction tensors are evaluated + against exactly the same BK the programs were learned on, rather than always the + plain ``atoms`` set. For the ``chebi_fg_rules`` / ``chebi_fg_learned_rules`` sets the + functional-group rule clauses are added to the BK directly (rather than pre-evaluated + into facts): the caller grounds and solves the combined program, which derives them. + + All work is scoped to ``rows``, so this can be called per molecule (e.g. to bound + Clingo grounding memory). ``aux_predicates`` (for ``llm_generated_fgs``) are the + name-deduplicated predicates gathered across all classes; their extensions are + evaluated on ``rows`` here. + + ``rule_dependencies`` (``llm_generated_rules``) are the library programs ``rule_programs`` + build on. They are ground alongside but emit no facts of their own. Pass them when the + caller has already resolved them — resolving here instead costs a full parse of the + library per call, and needs ``aux_library_dir`` to point at the right one. + """ + prolog_lines, _ = build_background_chemlog( + rows, aux_predicates=aux_predicates, aux_timeout=aux_timeout, aux_failures=aux_failures, + predicate_set=predicate_set, fowl_smarts=fowl_smarts, + ) + prolog_lines = list(prolog_lines) + + if predicate_set in ("chembl_fgs", "chebi_fgs"): + fg_data = get_chembl_fgs(rows) if predicate_set == "chembl_fgs" else get_chebi_fgs(rows) + fg_lines, _ = build_background_fg_data(fg_data, rows, source=predicate_set) + prolog_lines += fg_lines + + if predicate_set in ("chebi_fg_rules", "chebi_fg_learned_rules"): + rule_lines, _ = build_background_chebi_fg_rules( + CHEBI_FG_RULES_PATH if predicate_set == "chebi_fg_rules" else CHEBI_FG_LEARNED_RULES_PATH + ) + prolog_lines += rule_lines + + # llm_generated_rules: recompute the class's aux_* extensions exactly as build_bk + # does (ground each rule over atom + computed facts) and append them as facts, so a + # learned program's aux_* body literals resolve. Computed facts stay local to the + # grounding and are not added to the returned BK. + if predicate_set == "llm_generated_rules" and rule_programs: + eval_facts = list(prolog_lines) + if computed_facts: + eval_facts += build_computed_facts(rows) + mol_ids = [str(i) for i in rows.index] + if rule_dependencies is None: + rule_dependencies = resolve_rule_dependencies(rule_programs, aux_library_dir) + try: + extensions = derive_rule_extensions( + rule_programs + rule_dependencies, eval_facts, mol_ids, + ) + except (RuntimeError, MemoryError) as e: + print(f"Grounding failed ({e}); returning background knowledge without aux_* facts.") + extensions = {} + for rp in rule_programs: + emitted = set() + for arg_tuples in extensions.get(rp.name, {}).values(): + for args in arg_tuples: + line = f"{rp.name}({','.join(args)})." + if line not in emitted: + emitted.add(line) + prolog_lines.append(line) + + return prolog_lines + + +def build_background_chebi_fg_rules(rules_path=None): + """Load ChEBI functional group rules from a Prolog file and return them as BK lines and body predicates. + + Each rule defines a chebi_XXXXX(M) predicate in terms of atom-level predicates. + These are added as Prolog rules to the BK and as body_pred entries (arity 1) in the bias. + """ + if rules_path is None: + rules_path = CHEBI_FG_RULES_PATH + + prolog_lines = [f"% ChEBI FG rules from {os.path.basename(rules_path)}"] + body_predicates = [] + seen_predicates = set() + + with open(rules_path, "r") as f: + for line in f: + line = line.strip() + if not line or line.startswith("%"): + continue + prolog_lines.append(line) + # Extract predicate name from head: chebi_XXXXX(M) :- ... + pred_name = line.split("(")[0].strip() + if pred_name and pred_name not in seen_predicates: + seen_predicates.add(pred_name) + body_predicates.append(pred_name) + + print(f"Loaded {len(body_predicates)} ChEBI FG rule predicates from {rules_path}") + return prolog_lines, body_predicates + + +def build_background_fg_data(fg_data: dict[int, list[str]], rows, source: Literal["chembl_fgs", "chebi_fgs"]): + lines_by_predicate = dict() + + for row in rows.itertuples(): + if row.Index not in fg_data: + print(f"Warning: No functional group data found for CHEBI:{row.Index} in source {source}. This molecule will only have atom and bond predicates in the background knowledge.") + continue + for fg in fg_data[row.Index]: + if fg not in lines_by_predicate: + lines_by_predicate[fg] = [] + lines_by_predicate[fg].append(f"{fg}({row.Index}).") + total_lines = [line for lines in lines_by_predicate.values() for line in lines] + return total_lines, [(pred, 1) for pred in lines_by_predicate.keys()] + + +if __name__ == "__main__": + builder = ILPProblemBuilder( + chebi_version=251, + predicate_set="atoms", + ) + target_ids = ["134362"] + builder.build_examples(target_ids) diff --git a/chebILP/molecule_processing/prepare_chemlog_preds.py b/chebILP/molecule_processing/prepare_chemlog_preds.py new file mode 100644 index 0000000..7a9d590 --- /dev/null +++ b/chebILP/molecule_processing/prepare_chemlog_preds.py @@ -0,0 +1,217 @@ +""" +ChemLog predictions for the validation and test splits, in the tensor format +``ensemble_eval.load_dl_preds()`` reads. + + python -m chebILP.molecule_processing.prepare_chemlog_preds \\ + --chebi_version 251 \\ + --output_dir data/chemlog_v251 + +Produces, per split: + _preds_chemlog.npy float32 (n_molecules x n_classes) + _preds_chemlog_metadata.json {"mol_order": [...], "class_labels": [...]} + +ChemLog is the rule-based model python-chebifier contributes to its ensemble +(``chebifier.prediction_models.chemlog_predictor.ChemlogAllPredictor``). That class +cannot be imported here — chebifier's ``prediction_models`` package pulls in the neural +predictors — so the same three sub-classifiers are driven directly: + + peptides chemlog, 'algo' strategy, over 16 peptide classes + x molecular entity chemlog_extra, one class per chemical element present + organo-x compound chemlog_extra, one class per element bonded to carbon + +Columns are the classes a ChemLog rule decides directly, so every cell is a 1.0 or a +0.0 the rules stand behind. NaN appears only for a molecule ChemLog could not read. +Unlike chebifier's ensemble member, no positive is propagated up the ChEBI hierarchy: +a superclass no rule decides gets no column. +""" + +import argparse +import json +import os + +import networkx as nx +import numpy as np +from rdkit import Chem + +# The classes chemlog's peptide classifier decides, from ChemlogPeptidesPredictor minus +# 25696 and 25697: chemlog's resolve_chebi_classes emits 22563 / 36961 for those charge +# cases, so no rule can ever put a molecule in them. +# fmt: off +PEPTIDE_LABELS = [ + "15841", "16670", "24866", "25676", "27369", "46761", "47923", "48030", "48545", + "60194", "60334", "60466", "64372", "65061", "90799", "155837", +] +# fmt: on + + +def to_mol(molecule): + """Kekulise a copy of the molecule, as chemlog's classifiers expect. Mirrors + chebifier.utils.to_mol, which keeps a molecule that fails to kekulise.""" + if molecule is None: + return None + molecule = Chem.Mol(molecule) + try: + Chem.Kekulize(molecule) + except Chem.KekulizeException as e: + print(f"Failed to kekulise {Chem.MolToSmiles(molecule)}: {e}") + return molecule + + +class ChemLogPredictor: + """The three chemlog classifiers of chebifier's ``chemlog`` ensemble member, + predicting a ``{chebi_id: 0/1}`` dict per molecule.""" + + def __init__(self, chebi_graph: nx.DiGraph, chebi_version: int): + from chemlog.cli import CLASSIFIERS + from chemlog_extra.alg_classification.by_element_classification import ( + OrganoXCompoundClassifier, + XMolecularEntityClassifier, + ) + + # Both classifiers derive their element-to-class mapping from the ChEBI graph and + # cache it under data/chebi_v/ relative to the working directory. + self.element_classifiers = [ + XMolecularEntityClassifier(chebi_graph=chebi_graph, chebi_version=chebi_version), + OrganoXCompoundClassifier(chebi_graph=chebi_graph, chebi_version=chebi_version), + ] + self.peptide_classifiers = {key: cls() for key, cls in CLASSIFIERS["algo"].items()} + + @property + def classes(self) -> list[str]: + """Every class the rules decide, so that the column layout does not depend on + which classes a particular split happens to hit.""" + decided = set(PEPTIDE_LABELS) + for classifier in self.element_classifiers: + decided.update(classifier.element_class_mapping.values()) + return sorted(decided, key=int) + + def predict_peptides(self, mol) -> dict[str, int]: + from chemlog.cli import strategy_call + + predicted = strategy_call("algo", self.peptide_classifiers, mol)["chebi_classes"] + return {label: int(label in predicted) for label in PEPTIDE_LABELS} + + def predict_elements(self, mols: list) -> list[dict[str, int]]: + merged: list[dict[str, int]] = [dict() for _ in mols] + for classifier in self.element_classifiers: + for row, result in zip(merged, classifier.classify(list(mols))): + row.update({cls: int(hit) for cls, hit in result.items()}) + return merged + + def predict(self, mols: list) -> list[dict[str, int] | None]: + """One dict per molecule, None where the molecule is missing.""" + kekulized = [to_mol(mol) for mol in mols] + usable = [i for i, mol in enumerate(kekulized) if mol is not None] + + results: list[dict[str, int] | None] = [None] * len(mols) + element_preds = self.predict_elements([kekulized[i] for i in usable]) + for i, elements in zip(usable, element_preds): + results[i] = elements + for i in usable: + results[i].update(self.predict_peptides(kekulized[i])) + return results + + def on_finish(self): + for classifier in self.peptide_classifiers.values(): + classifier.on_finish() + + +def to_dense(predictions: list, classes: list[str]) -> np.ndarray: + """Stack the per-molecule dicts into a float32 matrix over *classes*, NaN where + ChemLog made no statement.""" + col_of = {cls: idx for idx, cls in enumerate(classes)} + + scores = np.full((len(predictions), len(classes)), np.nan, dtype="float32") + for row, pred in enumerate(predictions): + if not pred: + continue + for cls, value in pred.items(): + scores[row, col_of[cls]] = value + return scores + + +def prepare_chemlog_preds( + chebi_version: int, + output_dir: str, + splits: list[str], + base_dir: str = "data", + three_star_only: bool = True, + min_pos_samples: int = 25, + batch_size: int = 500, + limit: int | None = None, +): + import tqdm + + from chebILP.molecule_processing.data_preparation import ChEBIDataset + + dataset = ChEBIDataset( + chebi_version=chebi_version, + three_star_only=three_star_only, + base_dir=base_dir, + min_pos_samples=min_pos_samples, + ) + molecules = dataset.molecules + splits_df = dataset.load_splits_from_csv() + splits_df["id"] = splits_df["id"].astype(str) + + predictor = ChemLogPredictor(chebi_graph=dataset.chebi_graph, chebi_version=chebi_version) + + class_labels = predictor.classes + print(f"ChemLog decides {len(class_labels)} classes") + + os.makedirs(output_dir, exist_ok=True) + for split in splits: + split_ids = splits_df.loc[splits_df["split"] == split, "id"] + mol_order = [i for i in split_ids if i in molecules.index] + if limit: + mol_order = mol_order[:limit] + print(f"\n{split}: {len(mol_order)} molecules") + + predictions: list = [] + for start in tqdm.tqdm(range(0, len(mol_order), batch_size), desc=split): + batch = mol_order[start : start + batch_size] + predictions.extend(predictor.predict(molecules.loc[batch, "mol"].tolist())) + + scores = to_dense(predictions, class_labels) + npy_path = os.path.join(output_dir, f"{split}_preds_chemlog.npy") + meta_path = os.path.join(output_dir, f"{split}_preds_chemlog_metadata.json") + np.save(npy_path, scores) + with open(meta_path, "w") as f: + json.dump({"mol_order": mol_order, "class_labels": class_labels}, f, indent=2) + + print(f" Saved: {npy_path} (shape: {scores.shape}, {int((scores == 1).sum())} positive cells)") + print(f" Saved: {meta_path}") + + predictor.on_finish() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Produce ChemLog predictions for the validation and test splits." + ) + parser.add_argument("--chebi_version", "-v", type=int, default=251, help="ChEBI version.") + parser.add_argument("--base_dir", type=str, default="data", help="Dataset base directory.") + parser.add_argument("--include_two_star", "-2", action="store_true", + help="Include two-star molecules (default: three-star only).") + parser.add_argument("--min_pos_samples", type=int, default=25, + help="Minimum positive samples per class; selects the processed subdirectory.") + parser.add_argument("--output_dir", type=str, default=None, + help="Output directory (default: data/chemlog_v).") + parser.add_argument("--splits", type=str, nargs="+", default=["validation", "test"], + help="Splits to predict on.") + parser.add_argument("--batch_size", type=int, default=500, help="Molecules per progress step.") + parser.add_argument("--limit", type=int, default=None, + help="Only predict the first N molecules per split.") + + args = parser.parse_args() + + prepare_chemlog_preds( + chebi_version=args.chebi_version, + output_dir=args.output_dir or os.path.join("data", f"chemlog_v{args.chebi_version}"), + splits=args.splits, + base_dir=args.base_dir, + three_star_only=not args.include_two_star, + min_pos_samples=args.min_pos_samples, + batch_size=args.batch_size, + limit=args.limit, + ) diff --git a/chebILP/predicate_generation/auxiliary_generation.py b/chebILP/predicate_generation/auxiliary_generation.py index c8e4934..808a031 100644 --- a/chebILP/predicate_generation/auxiliary_generation.py +++ b/chebILP/predicate_generation/auxiliary_generation.py @@ -22,7 +22,7 @@ from pydantic import BaseModel, ConfigDict -from chebILP.predicate_generation.auxiliary_predicates import set_class_predicates +from chebILP.predicate_generation.auxiliary_predicates import load_class_map, set_class_predicates from chebILP.ilp_path_manager import get_aux_generation_log_path from chebILP.predicate_generation.llm_client import structured_completion from chebILP.utils import sort_labels_by_hierarchy @@ -108,7 +108,7 @@ def format_candidates(candidates: list[dict], with_kind: bool = False) -> str: def generate_one(prompt: str, system: str, model: str, selection_model, api_base: str | None = None): - """Ask the model for one class's selection. Returns ``(parsed, raw_json_text)``.""" + """Ask the model for one class's selection. Returns ``(parsed, raw_json_text, attempts)``.""" return structured_completion(model, system, prompt, selection_model, api_base=api_base) @@ -149,6 +149,16 @@ def to_source(self, item) -> str: def accept(self, source, label, ctx) -> tuple[bool, str]: """Validate a program. Returns ``(accepted, reason)``; reason is shown when rejected.""" + def accept_reused(self, stem, ctx) -> tuple[bool, str]: + """Validate a library program the model chose to reuse, same contract as :meth:`accept`. + + Retrieval matches a candidate on its name and description, which says nothing about how + it behaves on *this* class's molecules, so a reuse is worth the same check a new program + gets. Pipelines whose validation needs the class's molecules override this; the default + keeps every reuse, which is the behaviour of a pipeline that cannot tell. + """ + return True, "" + @abstractmethod def add_to_library(self, source, chebi_id): """Store the program. Returns ``(stem, saved)`` or ``None``.""" @@ -173,6 +183,15 @@ def prepare(self, blocks, ctx) -> None: def retriever_entry(self, stem, saved) -> dict: return {"name": saved.name, "description": saved.description, "kind": "rule", "stem": stem} + def rejection_code(self, label, ctx) -> str | None: + """Stable identifier for *why* ``label`` was rejected, or ``None`` if uncategorised. + + Recorded alongside the human-readable reason so a later pass — a repair prompt that + feeds the failure back to the model, say — can branch on the error type instead of + parsing prose. Only pipelines that categorise their rejections override this. + """ + return None + def describe(self, saved, stem, reason) -> str: return f"{saved.name} -> library/{stem} ({reason}): {saved.description}" @@ -185,12 +204,17 @@ def generate_for_class(self, chebi_id, info) -> int: candidates = self.retriever.retrieve(query, top_k=self.top_k) if len(self.retriever) else [] prompt = self.build_user_prompt(chebi_id, info, candidates, ctx) - parsed, raw = None, None + parsed, raw, attempts = None, None, [] try: - parsed, raw = generate_one(prompt, self.system_prompt, self.model, self.selection_model, self.api_base) + parsed, raw, attempts = generate_one( + prompt, self.system_prompt, self.model, self.selection_model, self.api_base + ) + except Exception as e: + attempts = getattr(e, "_chebilp_attempts", []) + raise finally: # Log the exchange even if the request failed; rewritten below once resolved. - log_path = self._write_log(chebi_id, info, prompt, parsed, raw, None) + log_path = self._write_log(chebi_id, info, prompt, parsed, raw, None, attempts) print(f" logged full exchange to {log_path}") reused_stems: list[str] = [] @@ -204,32 +228,55 @@ def generate_for_class(self, chebi_id, info) -> int: blocks = [(f"chebi_{chebi_id}_block_{i}", self.to_source(item)) for i, item in enumerate(parsed.new)] self.prepare(blocks, ctx) + rejected: list[dict] = [] + + # Reuses face the same gate as new programs. A reuse a new program builds on is still + # pulled back in by ``resolve_rule_dependencies`` at build_bk time, so dropping one here + # only keeps it out of the ILP's feature set — it never breaks a body that needs it. + kept_reused: list[dict] = [] + for stem in reused_stems: + ok, reason = self.accept_reused(stem, ctx) + if not ok: + code = self.rejection_code(stem, ctx) + print(f" rejected reused {stem}" + (f" [{code}]" if code else "") + f": {reason}") + record = {"name": stem, "reason": reason, "reused": True} + if code is not None: + record["code"] = code + rejected.append(record) + continue + kept_reused.append({"name": stem, "reason": reason}) + new_stems: list[str] = [] - rejected: list[str] = [] + new_records: list[dict] = [] seen: set[str] = set() for item, (label, source) in zip(parsed.new, blocks): ok, reason = self.accept(source, label, ctx) if not ok: - print(f" rejected {item.name}: {reason}") - rejected.append(f"{item.name} ({reason})") + code = self.rejection_code(label, ctx) + print(f" rejected {item.name}" + (f" [{code}]" if code else "") + f": {reason}") + record = {"name": item.name, "reason": reason} + if code is not None: + record["code"] = code + rejected.append(record) continue added = self.add_to_library(source, chebi_id) if added is None: - rejected.append(f"{item.name} (invalid program)") + rejected.append({"name": item.name, "reason": "invalid program"}) continue stem, saved = added if saved.name in seen: continue seen.add(saved.name) new_stems.append(stem) + new_records.append({"name": saved.name, "stem": stem, "reason": reason}) self.retriever.add_entry(self.retriever_entry(stem, saved)) print(f" new {self.describe(saved, stem, reason)}") - stems = reused_stems + new_stems + stems = [r["name"] for r in kept_reused] + new_stems set_class_predicates(chebi_id, stems, problem_dir=self.library_dir) - selection = {"reused": reused_stems, "new": new_stems, "rejected": rejected} - self._write_log(chebi_id, info, prompt, parsed, raw, selection) + selection = {"reused": kept_reused, "new": new_records, "rejected": rejected} + self._write_log(chebi_id, info, prompt, parsed, raw, selection, attempts) return len(stems) def run(self, chebi_graph, chebi_ids) -> int: @@ -238,8 +285,14 @@ def run(self, chebi_graph, chebi_ids) -> int: self.retriever = self.build_retriever() print(f" {len(self.retriever)} {self.noun}(s) in the library.") + done = set(load_class_map(self.library_dir)) + if done: + print(f" Resuming: {len(done)} class(es) already in class_map.json will be skipped.") + total = 0 for chebi_id in sort_labels_by_hierarchy(chebi_ids, chebi_graph): + if str(chebi_id) in done: + continue info = get_class_info(chebi_graph, chebi_id) print(f"CHEBI:{chebi_id} ({info['name']})...") try: @@ -253,27 +306,74 @@ def run(self, chebi_graph, chebi_ids) -> int: # --- logging ----------------------------------------------------------------- - def _write_log(self, chebi_id, info, prompt, parsed, raw, selection): + def _write_log(self, chebi_id, info, prompt, parsed, raw, selection, attempts=None): + attempts = attempts or [] + failed = [a for a in attempts if a.get("error")] + costs = [a["cost"] for a in attempts if a.get("cost") is not None] + total_cost = sum(costs) if costs else None log_path = get_aux_generation_log_path(chebi_id, base_dir=self.library_dir) with open(log_path, "w", encoding="utf-8") as f: f.write(f"# Auxiliary-{self.noun} generation log — CHEBI:{chebi_id} ({info['name']})\n\n") - f.write(f"- Timestamp: {time.strftime('%Y-%m-%d %H:%M:%S')}\n- Model: {self.model}\n\n") + f.write(f"- Timestamp: {time.strftime('%Y-%m-%d %H:%M:%S')}\n- Model: {self.model}\n") + if attempts: + cost_str = f"${total_cost:.4f}" if total_cost is not None else "unknown" + f.write(f"- Cost: {cost_str} across {len(attempts)} call(s), {len(failed)} reasked\n") + f.write("\n") if selection is not None: - f.write("## Resolved selection\n\n") - f.write(f"- Reused from library: {', '.join(selection['reused']) or '(none)'}\n") - f.write(f"- New {self.noun}s added: {', '.join(selection['new']) or '(none)'}\n") - f.write(f"- Rejected: {'; '.join(selection['rejected']) or '(none)'}\n\n") + f.write(self._format_selection(selection)) if parsed is not None: f.write(f"## Model reasoning\n\n{parsed.reasoning}\n\n") f.write(self._format_output(parsed)) f.write("## System prompt\n\n```\n" + self.system_prompt + "\n```\n\n") f.write("## User prompt\n\n```\n" + prompt + "\n```\n\n") - if parsed is None: + if parsed is None and not failed: f.write("## Raw LLM response\n\n```json\n") f.write(raw if raw is not None else "(no response — request failed)") - f.write("\n```\n") + f.write("\n```\n\n") + if failed: + # Kept at the bottom: the resolved result is what matters; the reasks are + # diagnostic context for why an extra call (or several) was needed. + f.write(self._format_failed_attempts(failed)) return log_path + def _format_selection(self, selection) -> str: + """Render the resolved selection: reused, accepted (with fire fraction), rejected.""" + parts = ["## Resolved selection\n\n"] + if selection["reused"]: + parts.append("- Reused from library:\n") + for r in selection["reused"]: + parts.append(f" - `{r['name']}`" + (f" — {r['reason']}\n" if r.get("reason") else "\n")) + else: + parts.append("- Reused from library: (none)\n") + if selection["new"]: + parts.append(f"- New {self.noun}s added:\n") + for r in selection["new"]: + parts.append(f" - `{r['name']}` — {r['reason']}\n") + else: + parts.append(f"- New {self.noun}s added: (none)\n") + if selection["rejected"]: + parts.append("- Rejected:\n") + for r in selection["rejected"]: + code = f"**[{r['code']}]** " if r.get("code") else "" + origin = " (reused)" if r.get("reused") else "" + parts.append(f" - `{r['name']}`{origin} — {code}{r['reason']}\n") + else: + parts.append("- Rejected: (none)\n") + parts.append("\n") + return "".join(parts) + + def _format_failed_attempts(self, failed) -> str: + """Render reasked calls (malformed / schema-invalid output) at the log's tail.""" + parts = [f"## Failed attempts ({len(failed)} reasked)\n\n"] + for i, a in enumerate(failed, 1): + cost = a.get("cost") + cost_str = f" — cost ${cost:.4f}" if cost is not None else "" + parts.append(f"### Attempt {i}{cost_str}\n\n") + parts.append(f"- Error: {a['error']}\n\n") + raw = a.get("raw") + parts.append("```json\n" + (raw if raw else "(no raw response captured)") + "\n```\n\n") + return "".join(parts) + def _format_output(self, parsed) -> str: """Render the model's parsed answer as readable Markdown.""" parts = ["## Model output\n\n"] diff --git a/chebILP/predicate_generation/auxiliary_rules.py b/chebILP/predicate_generation/auxiliary_rules.py index 5ddf4e5..514f081 100644 --- a/chebILP/predicate_generation/auxiliary_rules.py +++ b/chebILP/predicate_generation/auxiliary_rules.py @@ -14,9 +14,10 @@ validate-and-reject step and ``build_bk`` attribute an extension to molecules. A class's programs are grounded TOGETHER, so one program may use a predicate another defines -and helpers are shared. The flip side is that names are shared too: two programs of the same -class defining one name differently merge into a single extension. Programs of *different* -classes must therefore be name-qualified before they meet in one grounding. +and helpers are shared. Names are shared too, but that is safe across classes: a name in the +library backs exactly one program (:func:`add_rule_to_library` disambiguates content that +would collide with the class id), so programs of different classes can meet in one grounding +as they are. Storage reuses the shared-library layout of the Python pipeline, but with ``.pl`` files and a separate default directory so the two libraries never mix:: @@ -28,9 +29,9 @@ class defining one name differently merge into a single extension. Programs of * A rule program file carries a two-line header naming the predicate and describing it:: - % PREDICATE_NAME: aux_at_least_three_rings - % DESCRIPTION: molecule has at least three rings - aux_at_least_three_rings(M) :- N = #count{ S : ring_size(M,S) }, N >= 3. + % PREDICATE_NAME: aux_at_least_three_ring_atoms + % DESCRIPTION: molecule has at least three atoms lying in a ring + aux_at_least_three_ring_atoms(M) :- has_atom(M,_), 3 <= #count{ A : has_atom(M,A), in_ring(A) }. """ from __future__ import annotations @@ -48,6 +49,10 @@ class defining one name differently merge into a single extension. Programs of * # Rule library, kept separate from the Python-program library (data/llm_generated_predicates). DEFAULT_AUX_RULE_LIBRARY_DIR: str = os.path.join("data", "llm_generated_rules") +# Wall-clock ceiling for a whole ``derive_rule_extensions`` call, across all of its batches. +# Generous: it is a backstop against a program that never finishes, not a performance budget. +DEFAULT_GROUNDING_TIMEOUT: float = 300.0 + _NAME_RE = re.compile(r"^%\s*PREDICATE_NAME\s*:\s*(.+?)\s*$", re.IGNORECASE) _DESC_RE = re.compile(r"^%\s*DESCRIPTION\s*:\s*(.+?)\s*$", re.IGNORECASE) @@ -97,6 +102,7 @@ def parse_rule_program(source: str, source_file: str = "") -> RuleProgra _HAS_ATOM_RE = re.compile(r"^\s*has_atom\(\s*([^,\s]+)\s*,\s*([^)\s]+)\s*\)\s*\.\s*$") +_FACT_RE = re.compile(r"^\s*[a-z_][A-Za-z0-9_]*\s*\((.*)\)\s*\.\s*$") def atom_to_molecule(facts: list[str]) -> dict[str, str]: @@ -109,6 +115,36 @@ def atom_to_molecule(facts: list[str]) -> dict[str, str]: return mapping +def group_facts_by_molecule(facts: list[str], mol_ids) -> list[list[str]]: + """Split ``facts`` into one fact list per molecule, in ``mol_ids`` order. + + Which molecule a fact belongs to is resolved through the ``has_atom`` facts, not the + spelling of the atom id — the same rule :func:`derive_rule_extensions` uses to attribute + a derived atom. A fact naming neither an atom nor a known molecule (there are normally + none) is unattributable and is copied into *every* group, so an unrecognised fact shape + can only ever be over-supplied, never silently dropped. + """ + atom2mol = atom_to_molecule(facts) + known = {str(m) for m in mol_ids} + + by_mol: dict[str, list[str]] = {m: [] for m in known} + shared: list[str] = [] + for line in facts: + m = _FACT_RE.match(line) + args = [a.strip() for a in m.group(1).split(",")] if m else [] + owner = next( + (atom2mol[a] for a in args if a in atom2mol), + next((a for a in args if a in known), None), + ) + if owner is None: + shared.append(line) + elif owner in by_mol: + by_mol[owner].append(line) + + ordered = [str(m) for m in mol_ids if str(m) in by_mol] + return [by_mol[m] + shared for m in dict.fromkeys(ordered) if by_mol[m] or shared] + + def rule_program_error(source: str) -> str | None: """Clingo's complaint if ``source`` is not a well-formed program, else ``None``. @@ -119,6 +155,9 @@ def rule_program_error(source: str) -> str | None: """ import clingo + from chebILP.evaluation.clingo_eval import _patch_clingo_logger_decode + + _patch_clingo_logger_decode() messages: list[str] = [] ctl = clingo.Control(logger=lambda code, msg: messages.append(msg)) try: @@ -130,13 +169,292 @@ def rule_program_error(source: str) -> str | None: return None -def derive_rule_extensions(progs, facts: list[str], mol_ids, timeout: float | None = None) -> dict[str, dict[str, list[tuple[str, ...]]]]: +_COMMENT_RE = re.compile(r"%.*") +_PRED_RE = re.compile(r"\b([a-z_][A-Za-z0-9_]*)\s*\(") +# ``L = L0 + 1`` and friends: a variable assigned another variable offset by a constant. +_INCREMENT_RE = re.compile(r"\b([A-Z][A-Za-z0-9_]*)\s*=\s*([A-Z][A-Za-z0-9_]*)\s*[+-]\s*\d+") +# the same growth written straight into a head argument: ``aux_p(A, D, N+1)`` +_ARG_ARITH_RE = re.compile(r"\b([A-Z][A-Za-z0-9_]*)\s*[+-]\s*\d+") + + +def _clauses(source: str) -> list[tuple[str, str]]: + """``(head, body)`` for every clause in ``source``; a fact gets an empty body.""" + parts = re.split(r"\.(?=\s|$)", _COMMENT_RE.sub("", source)) + clauses = [] + for chunk in parts: + chunk = chunk.strip() + if chunk: + head, _, body = chunk.partition(":-") + clauses.append((head.strip(), body.strip())) + return clauses + + +def _recursive_predicates(progs) -> set[str]: + """Predicate names that depend on themselves, directly or through other clauses.""" + edges: dict[str, set[str]] = {} + for prog in progs: + for head, body in _clauses(prog.source): + m = _PRED_RE.match(head) + if m is not None and body: + edges.setdefault(m.group(1), set()).update(_PRED_RE.findall(body)) + + reach = {name: set(deps) for name, deps in edges.items()} + changed = True + while changed: # transitive closure; a class's dependency graph is a handful of nodes + changed = False + for name, deps in reach.items(): + grown = deps.union(*(reach.get(d, set()) for d in deps)) if deps else deps + if grown != deps: + reach[name] = grown + changed = True + return {name for name, deps in reach.items() if name in deps} + + +def _bounded_above(clause: str, variables: set[str]) -> bool: + """Whether any of ``variables`` is capped by an integer constant in ``clause``.""" + return any( + re.search(rf"\b{re.escape(v)}\s*(?:<=|<)\s*\d+", clause) + or re.search(rf"\b\d+\s*(?:>=|>)\s*{re.escape(v)}\b", clause) + for v in variables + ) + + +def unbounded_recursion_errors(progs) -> dict[str, str]: + """Programs whose grounding would never terminate, keyed by program name. + + A recursive clause that increments a counter derives a *new* ground atom at every step, + so its fixpoint is infinite and clingo grounds until memory runs out. The molecule graph + is always cyclic here — ``has_bond_to`` is symmetric, so a walk can step A -> B -> A even + in an acyclic molecule — and the counter keeps each lap distinct. Dropping the counter + makes the very same rule saturate, which is why plain reachability is safe and this is + not. Only a constant upper bound in the SAME clause stops the grounder; one applied by a + consumer clause comes too late. + + ``progs`` is inspected as a set, since the cycle may run through a sibling or a reused + library program, but only the clause carrying the increment is blamed. + """ + recursive = _recursive_predicates(progs) + errors: dict[str, str] = {} + for prog in progs: + for head, body in _clauses(prog.source): + m = _PRED_RE.match(head) + if m is None or not body or m.group(1) not in recursive: + continue + clause = f"{head} :- {body}" + growing = {v for pair in _INCREMENT_RE.findall(clause) for v in pair} + growing |= set(_ARG_ARITH_RE.findall(head[head.find("(") + 1:] if "(" in head else "")) + if growing and not _bounded_above(clause, growing): + errors[prog.name] = ( + f"non-terminating recursion: {m.group(1)} recurses while incrementing " + f"{'/'.join(sorted(growing))}, with no constant upper bound in that clause" + ) + break + return errors + + +_VAR_RE = re.compile(r"\b([A-Z][A-Za-z0-9_]*)\b") +_NOT_RE = re.compile(r"^not\s+") +_AGG_BRACES_RE = re.compile(r"\{[^{}]*\}") +# A literal with no predicate call: ``A != C``, ``R >= 5``, ``L = L0 + 1``, ``M = M``. +_COMPARISON_OPS_RE = re.compile(r"!=|<=|>=|<|>|=") + + +def _body_literals(body: str) -> list[str]: + """Top-level comma-separated literals of ``body``; commas inside ``(``/``{`` don't split.""" + literals, depth, current = [], 0, "" + for ch in body: + if ch in "({[": + depth += 1 + elif ch in ")}]": + depth -= 1 + if ch == "," and depth == 0: + literals.append(current) + current = "" + else: + current += ch + literals.append(current) + return [lit.strip() for lit in literals if lit.strip()] + + +class _Union: + """Minimal union-find over variable names.""" + + def __init__(self): + self.parent: dict[str, str] = {} + + def add(self, var: str) -> None: + self.parent.setdefault(var, var) + + def find(self, var: str) -> str: + while self.parent[var] != var: + self.parent[var] = self.parent[self.parent[var]] + var = self.parent[var] + return var + + def union(self, vars_) -> None: + vars_ = [v for v in vars_ if v in self.parent] + for other in vars_[1:]: + a, b = self.find(vars_[0]), self.find(other) + if a != b: + self.parent[a] = b + + def components(self) -> dict[str, set[str]]: + groups: dict[str, set[str]] = {} + for var in self.parent: + groups.setdefault(self.find(var), set()).add(var) + return groups + + +def _clause_components(head: str, body: str) -> tuple[dict[str, set[str]], list[str], set[str]]: + """``(variable components, negative literals, aggregate result variables)`` for a clause. + + Variables are joined only by *positive* predicate literals and by plain ``=``. A negated + literal cannot bind anything, so it does not join; nor does an inequality. Inside an + aggregate the element variables are local, but a variable that also occurs outside it is + global — so an aggregate joins exactly those of its variables that appear elsewhere. That + is what separates a molecule-scoped count from an unscoped one: ``#count{ A : has_atom(M,A), + … }`` mentions ``M`` inside and so ties its result to ``M``, while ``#count{ A : p(A) }`` + leaves the result standing alone. + """ + literals = _body_literals(body) + union = _Union() + for var in _VAR_RE.findall(f"{head} {body}"): + union.add(var) + + aggregate_results: set[str] = set() + for i, literal in enumerate(literals): + if _NOT_RE.match(literal): + continue + if "#" in literal: + elsewhere = set(_VAR_RE.findall(head)) + elsewhere.update(v for j, o in enumerate(literals) if j != i for v in _VAR_RE.findall(o)) + union.union([v for v in _VAR_RE.findall(literal) if v in elsewhere]) + # variables outside the braces are the aggregate's result, bound by it + outer = _VAR_RE.findall(_AGG_BRACES_RE.sub(" ", literal)) + aggregate_results.update(outer) + union.union(outer) + elif "(" in literal: + union.union(_VAR_RE.findall(literal)) + elif _COMPARISON_OPS_RE.search(literal): + if re.fullmatch(r"[^!<>=]*=[^=]*", literal): # plain equality does bind + union.union(_VAR_RE.findall(literal)) + + return union.components(), [lit for lit in literals if _NOT_RE.match(lit)], aggregate_results + + +# Stable identifiers for the static rejections, so a caller — a repair prompt for a second +# LLM pass, say — can branch on the error type instead of parsing the message. +ERROR_UNBOUNDED_RECURSION = "unbounded_recursion" +ERROR_CROSS_PRODUCT = "cross_product" +ERROR_NEGATION_SCOPE = "negation_scope" +ERROR_UNSCOPED_AGGREGATE = "unscoped_aggregate" + + +def cross_product_errors(progs) -> dict[str, tuple[str, str]]: + """Programs with an unjoined-variable clause, as ``{name: (code, message)}``. + + Such a clause's body splits into groups of variables that no positive literal connects, + so clingo grounds their **cartesian product** — over every atom in the instance, not just + the molecule's. Two harms: the extension is wrong (a molecule-level head fires because of + a *different* molecule's atoms) and the cost is quadratic in the molecules ground + together, which is what exhausted memory in ``build_bk``. Batching bounds the cost + (see :func:`derive_rule_extensions`); only rejection fixes the meaning. + + Three codes are reported, most specific first: + + - :data:`ERROR_UNSCOPED_AGGREGATE` — ``N = #count{ A : p(A) }`` with no ``has_atom(M,A)`` + inside, so every molecule is assigned the same instance-wide total. Cheap to ground and + therefore invisible to any memory guard; only wrong. + - :data:`ERROR_NEGATION_SCOPE` — ``not p(X), q(X)`` written for ``not (p(X), q(X))``, which + ``_SYSTEM_PROMPT`` already forbids. The clause stays *safe*, so clingo accepts it and + only this check catches it. + - :data:`ERROR_CROSS_PRODUCT` — anything else, normally a head variable never joined to the + atoms the body tests. + + The last two overlap where a pair predicate leaves its two atoms unjoined; the message + names the offending variables either way. + + Unlike :func:`unbounded_recursion_errors` this is per-program: a cross product is local to + one clause and cannot be closed by a sibling. + """ + errors: dict[str, tuple[str, str]] = {} + for prog in progs: + for head, body in _clauses(prog.source): + if not body or _PRED_RE.match(head) is None: + continue + components, negatives, aggregate_results = _clause_components(head, body) + # A variable no positive literal mentions is not a component of its own; only + # groups that actually generate atoms count. + positive_vars = { + v + for lit in _body_literals(body) + if not _NOT_RE.match(lit) and "(" in lit + for v in _VAR_RE.findall(_AGG_BRACES_RE.sub(" ", lit) if "#" in lit else lit) + } + generating = [c for c in components.values() if c & positive_vars] + if len(generating) < 2: + continue + + spanned = next( + ( + lit + for lit in negatives + if sum(1 for c in generating if c & set(_VAR_RE.findall(lit))) > 1 + ), + None, + ) + unscoped = next((c for c in generating if c <= aggregate_results), None) + groups = " and ".join("{" + ", ".join(sorted(c)) + "}" for c in generating) + if unscoped is not None: + errors[prog.name] = ( + ERROR_UNSCOPED_AGGREGATE, + f"aggregate not scoped to the molecule: {', '.join(sorted(unscoped))} is " + f"counted over every molecule at once, so each molecule gets the same total. " + f"Bind the element to the molecule inside the aggregate, as in " + f"#count{{ A : has_atom(M,A), ... }}", + ) + elif spanned is not None: + errors[prog.name] = ( + ERROR_NEGATION_SCOPE, + f"negation over a conjunction: `{spanned.strip()}` is the only thing linking " + f"{groups}, but a negated literal cannot bind. Define a helper predicate for " + f"the whole conjunction and negate that helper instead", + ) + else: + errors[prog.name] = ( + ERROR_CROSS_PRODUCT, + f"unjoined variables: {groups} are never linked by a positive literal, so " + f"this clause grounds as their cross product over every molecule. Join them " + f"— usually via has_atom(M,...) for each atom variable", + ) + break + return errors + + +def static_rule_errors(progs) -> dict[str, tuple[str, str]]: + """All static rejections for ``progs``, as ``{program name: (code, message)}``. + + The single entry point for checks that must run *before* clingo sees a program, because + the failure they describe is one grounding cannot survive (non-termination) or one it + cannot detect (a cross product grounds "successfully", just wrongly and enormously). + """ + errors = { + name: (ERROR_UNBOUNDED_RECURSION, message) + for name, message in unbounded_recursion_errors(progs).items() + } + for name, entry in cross_product_errors(progs).items(): + errors.setdefault(name, entry) + return errors + + +def derive_rule_extensions(progs, facts: list[str], mol_ids, timeout: float | None = DEFAULT_GROUNDING_TIMEOUT, batch_size: int | None = None) -> dict[str, dict[str, list[tuple[str, ...]]]]: """Ground ``progs`` together and group each one's derived atoms by molecule. All programs go into a single clingo instance, so one may build on predicates another defines. Names are consequently shared across ``progs``: two programs that define the same - predicate differently contribute to a single extension. Within one class that is the - intent; across classes the names must be qualified first (see ``test._qualify_aux_rules``). + predicate differently would contribute to a single extension. The library rules that out + by construction — one name, one program — so programs from several classes may be passed + together. A head may have any arity: a molecule predicate ``aux_x(M)``, an atom one ``aux_x(A)``, a pair ``aux_x(A1,A2)``, or any mix. A derived atom belongs to a molecule when one of its @@ -144,12 +462,30 @@ def derive_rule_extensions(progs, facts: list[str], mol_ids, timeout: float | No rather than the spelling of the atom id. Arguments that are neither (plain numbers, say) attach the atom to nothing on their own. + The facts are split per molecule and ground in batches of ``batch_size`` rather than all + at once. For a well-formed program this changes nothing — every variable is joined to its + molecule, so no clause could span two of them anyway — but it bounds what a clause with + *unjoined* variables costs: such a clause grounds as a cross product over every atom in + the instance, which at class scale (650+ molecules) reaches 10^8 ground atoms and exhausts + memory. See :data:`~chebILP.evaluation.clingo_eval.GROUNDING_BATCH_SIZE`. + + ``timeout`` is the wall-clock ceiling for the whole call. + Returns ``{program_name: {mol_id: [arg_tuple, ...]}}``, with one entry per program. """ - from chebILP.evaluation.clingo_eval import ground_extensions + from chebILP.evaluation.clingo_eval import GROUNDING_BATCH_SIZE, ground_extensions_isolated + if batch_size is None: + batch_size = GROUNDING_BATCH_SIZE names = [p.name for p in progs] - derived = ground_extensions([p.source for p in progs], facts, names, timeout=timeout) + per_molecule = group_facts_by_molecule(facts, mol_ids) + fact_groups = [ + [line for group in per_molecule[i:i + batch_size] for line in group] + for i in range(0, len(per_molecule), batch_size) + ] or [facts] + derived = ground_extensions_isolated( + [p.source for p in progs], fact_groups, names, timeout=timeout, total_timeout=timeout, + ) atom2mol = atom_to_molecule(facts) known = {str(m) for m in mol_ids} @@ -242,6 +578,108 @@ def load_library_rules(base_dir: str) -> list[RuleProgram]: return programs +def _defined_predicates(progs) -> set[str]: + """Predicate names appearing in the head of any clause of ``progs``.""" + return { + m.group(1) + for prog in progs + for head, _ in _clauses(prog.source) + for m in [_PRED_RE.match(head)] + if m is not None + } + + +def _referenced_predicates(prog) -> set[str]: + """Predicate names appearing in the body of any clause of ``prog``.""" + return {name for _, body in _clauses(prog.source) for name in _PRED_RE.findall(body)} + + +_DEFINITION_INDEX_CACHE: dict[str, dict[str, list[str]]] = {} + + +def _definition_index(library_dir: str) -> dict[str, list[str]]: + """``{predicate: [program stem, ...]}`` over every clause head in the library. + + Memoised per directory: the index costs a full parse of the library (1600+ files) and + ``build_bk`` resolves against it once per class. + """ + if library_dir in _DEFINITION_INDEX_CACHE: + return _DEFINITION_INDEX_CACHE[library_dir] + + programs_dir = get_aux_programs_dir(base_dir=library_dir) + index: dict[str, list[str]] = {} + for fname in sorted(os.listdir(programs_dir)): + if not fname.endswith(".pl") or fname.startswith("_"): + continue + with open(os.path.join(programs_dir, fname), "r", encoding="utf-8") as f: + source = f.read() + stem = fname[:-3] + for head, _ in _clauses(source): + m = _PRED_RE.match(head) + if m is not None and stem not in index.setdefault(m.group(1), []): + index[m.group(1)].append(stem) + _DEFINITION_INDEX_CACHE[library_dir] = index + return index + + +def resolve_rule_dependencies(progs, library_dir: str | None = None) -> list[RuleProgram]: + """Library programs defining ``aux_*`` predicates ``progs`` reference but do not define. + + ``class_map.json`` records only the predicates a class *chose*, never the ones those + programs build on — the layering the generator asks for is invisible to it. So a program + like ``aux_amino_nitrogen(N) :- aux_organic_amine_nitrogen(N), not aux_n_acylated_amino(N).`` + loads with both dependencies missing, derives nothing, and silently costs the class a + predicate. Resolving them here fixes existing libraries without regenerating anything. + + The closure is transitive and cycle-safe. A predicate defined by several programs is + resolved to the one whose file *is* that predicate, else one already loaded, else the + first stem alphabetically (with a warning) — the library has no other notion of which + definition is canonical. + + Returns only the programs that had to be added; the caller keeps emitting extensions for + its own programs alone. + """ + library_dir = library_dir or DEFAULT_AUX_RULE_LIBRARY_DIR + index = _definition_index(library_dir) + + resolved: list[RuleProgram] = [] + defined = _defined_predicates(progs) + loaded_stems: set[str] = set() + pending = list(progs) + + while pending: + for name in _referenced_predicates(pending.pop()): + if not name.startswith("aux_") or name in defined: + continue + stems = index.get(name) + if not stems: + continue # nothing in the library defines it; it just grounds as empty + if name in stems: + stem = name # the program whose file IS this predicate is the canonical one + else: + stem = next((s for s in stems if s in loaded_stems), stems[0]) + if len(stems) > 1: + logger.warning( + "Auxiliary predicate %s is defined by %d library programs; using %s.", + name, len(stems), stem, + ) + if stem in loaded_stems: + defined.add(name) # already pulled in under a different predicate's name + continue + loaded_stems.add(stem) + path = aux_rule_path(stem, library_dir) + if not os.path.exists(path): + continue + with open(path, "r", encoding="utf-8") as f: + dep = parse_rule_program(f.read(), source_file=path) + if dep is None: + continue + resolved.append(dep) + pending.append(dep) + defined |= _defined_predicates([dep]) + return resolved + + def load_class_rules(chebi_id, library_dir: str | None = None) -> list[RuleProgram]: """Load the rule programs a ChEBI class uses (from ``class_map.json``). diff --git a/chebILP/predicate_generation/generate_auxiliary_predicates.py b/chebILP/predicate_generation/generate_auxiliary_predicates.py index 08f57de..a0daf27 100644 --- a/chebILP/predicate_generation/generate_auxiliary_predicates.py +++ b/chebILP/predicate_generation/generate_auxiliary_predicates.py @@ -54,7 +54,8 @@ - bSINGLE/bDOUBLE/bTRIPLE/bAROMATIC(Atom1, Atom2): bond type - bSTEREOCIS/bSTEREOTRANS(Atom1, Atom2): cis/trans bond stereochemistry - has_bond_to(Atom1, Atom2): any bond between two atoms -- in_ring(Atom), in_ringN(Atom), ringN(A1..AN): ring membership / N-membered rings (N up to 8) +- in_ring(Atom), in_ringN(Atom): ring membership / membership of an N-membered ring, for + every N that occurs. There is no N-ary ringN(A1..AN) naming a ring's atoms in order. - net_charge_positive/negative/neutral(Molecule), aromatic/aliphatic(Molecule): molecule-level - steroid_1..steroid_17(Atom): atom at a steroid-nucleus position\ """ diff --git a/chebILP/predicate_generation/generate_auxiliary_rules.py b/chebILP/predicate_generation/generate_auxiliary_rules.py index f27c19c..41adbfa 100644 --- a/chebILP/predicate_generation/generate_auxiliary_rules.py +++ b/chebILP/predicate_generation/generate_auxiliary_rules.py @@ -29,11 +29,13 @@ ) from chebILP.predicate_generation.auxiliary_rules import ( DEFAULT_AUX_RULE_LIBRARY_DIR, + ERROR_UNBOUNDED_RECURSION, add_rule_to_library, aux_rule_path, derive_rule_extensions, parse_rule_program, rule_program_error, + static_rule_errors, ) from chebILP.ilp_path_manager import get_exs_path from chebILP.utils import get_atom_id @@ -42,16 +44,16 @@ # Background facts the ILP system already provides. The rules may use any of these. _EXISTING_PREDICATES = """\ -- has_atom(M, A): molecule M contains atom A (bind atom variables through this) -- c(A), n(A), o(A), s(A), p(A), cl(A), br(A), f(A), i(A), se(A), ...: atom element (lowercase) -- charge0/charge_p/charge_n/charge1/charge_m1(A): formal charge of atom A +- has_atom(M, A): molecule M contains atom A +- c(A), n(A), p(A), cl(A), ...: atom element (lowercase) +- charge_[p|n|0|1|-1|...](A): charge of atom A (p = positive, n = negative, 0 = neutral, m1 = -1, etc.) - has_X_hs(A) / has_at_least_X_hs(A): attached-hydrogen counts -- cip_code_R(A), cip_code_S(A): R/S CIP stereochemistry (often ABSENT — do not over-rely) +- cip_code_R(A), cip_code_S(A): R/S CIP stereochemistry - bSINGLE/bDOUBLE/bTRIPLE/bAROMATIC(A1, A2): bond type (symmetric) - bSTEREOCIS/bSTEREOTRANS(A1, A2): cis/trans bond stereochemistry - has_bond_to(A1, A2): any bond between two atoms (symmetric) - in_ring(A): A is in a ring of any size -- in_ringN(A), ringN(A1..AN): ring membership / N-membered ring, ONLY for N up to 8 +- in_ringN(A): A is in an N-membered ring. - steroid_1..steroid_17(A): atom at a steroid-nucleus position - net_charge_positive/negative/neutral(M), aromatic/aliphatic(M): molecule-level\ """ @@ -59,8 +61,7 @@ # Extra molecule-level facts available only in this mode (never written to bk.pl themselves). _COMPUTED_PREDICATES = """\ - mol_weight(M, W): integer molecular weight of M (rounded), for weight thresholds -- ring_size(M, S): one fact per ring of M giving its size S — use this for rings LARGER than - 8 (which ringN/in_ringN cannot express) and to count rings via #count\ +- ring_size(M, S): M has a ring of size S.\ """ @@ -69,8 +70,7 @@ Logic Programming (ILP). Your task is to write AUXILIARY PREDICATES that help an ILP system distinguish a ChEBI chemical class from other molecules. -Each predicate has three parts, which you supply as separate fields. The program is a -logic program that needs to be parsed by clingo. +Each predicate has three parts. The program needs to be parsed by clingo. name: aux_snake_case_name description: @@ -79,94 +79,35 @@ Contract: - "program" holds ONLY clauses. Do not repeat the name or description inside it as comments. -- The head may take WHATEVER arguments suit the property, and its name MUST match the "name" - field. A predicate can describe a property of the molecule, an atom or a relation between atoms. -- Your programs for this class are grounded TOGETHER, so a program MAY use a predicate that - another one defines (including one you are reusing from the library) — build them up in - layers rather than repeating clauses. Define each helper ONCE, in the - program it most belongs to. Keep each program simple to avoid errors and make it easier to reuse later. - You MAY use clingo aggregates (#count, #sum), comparisons (=, !=, <, <=, >, >=) and negation-as-failure (not ...). You MAY define recursive helper predicates. -- Every variable in the head must be bound by the body. Never leave the program unsafe. -- There is no `or` and no parentheses for grouping in a body. Write one clause per - alternative — the head holds if ANY clause does (see aux_carboxyl_carbon below). - Don't use `a ; b` as an `or`. That is a syntax error. +- Don't use `a ; b` as an `or`. That is a syntax error. - A variable that appears ONLY inside an aggregate is LOCAL to it and does NOT bind the head. Bind it in the body outside the aggregate first: aux_x(M) :- has_atom(M,_), 2 = #count{ A : has_atom(M,A), aux_y(A) }. % M bound: OK aux_x(M) :- 2 = #count{ A : has_atom(M,A), aux_y(A) }. % M UNSAFE: rejected -- Negation-as-failure applies to ONE literal: `not p(X)` is valid, `not (p(X), q(X))` is NOT. - To negate a conjunction, define a helper predicate for it and negate that helper. - -There is a SHARED LIBRARY of auxiliary rules already written for other classes. You will be -shown the most relevant existing ones; prefer to REUSE those that fit over writing near-duplicates. - -Here a some examples of different kinds of useful predicates. - -A substructure or local pattern. Molecule-level, "does the molecule contain this?": - name: aux_has_azetidine_ring - description: contains a 4-membered ring holding a nitrogen - program: aux_has_azetidine_ring(M) :- has_atom(M,A), n(A), ring4(A,B,C,D). - -The same kind of pattern, but naming the ATOMS that match, so ILP can reason about where they -sit and join them to other atom predicates. Prefer this when the location matters: - name: aux_carboxyl_carbon - description: carbons that are the carbon of a carboxyl group (C=O with an -OH or -O anion) - program: aux_carboxyl_carbon(C) :- c(C), bDOUBLE(C,O1), o(O1), bSINGLE(C,O2), o(O2), has_1_hs(O2), O1 != O2. - aux_carboxyl_carbon(C) :- c(C), bDOUBLE(C,O1), o(O1), bSINGLE(C,O2), o(O2), charge_m1(O2), O1 != O2. - - -An atom PAIR, when the property is a relationship between two atoms: - name: aux_amide_bond - description: pairs of (carbonyl carbon, amide nitrogen) joined by an amide bond - program: aux_amide_bond(C,N) :- c(C), o(O), - bDOUBLE(C,O), has_bond_to(C,N), n(N). - -A count, expressed with aggregates. Use this if it is relevant HOW MANY of a group there are (N -carbons, N sugar units, one vs two carboxyls). - name: aux_exactly_two_ether_oxygens - description: has exactly two ether oxygens (O bonded to two carbons, no hydrogen) - program: aux_ether_oxygen(O) :- o(O), has_0_hs(O), has_bond_to(O,C1), - c(C1), has_bond_to(O,C2), c(C2), C1 != C2. - aux_exactly_two_ether_oxygens(M) :- has_atom(M,_), - 2 = #count{ O : has_atom(M,O), aux_ether_oxygen(O) }. - -An absence, via negation. Note it BUILDS ON aux_carboxyl_carbon above instead of restating -it — that is the layering to aim for. - name: aux_no_carboxyl - description: has no carboxyl group - program: aux_has_carboxyl(M) :- has_atom(M,C), aux_carboxyl_carbon(C). - aux_no_carboxyl(M) :- has_atom(M,_), not aux_has_carboxyl(M). - -A chain length, path or connectivity property, expressed with recursion: - name: aux_large_carbon_skeleton - description: has a connected carbon subgraph of at least 22 carbons - program: aux_carbon_reachable(A,A) :- c(A). - aux_carbon_reachable(A,D) :- aux_carbon_reachable(A,B), has_bond_to(B,D), c(D). - aux_large_carbon_skeleton(M) :- has_atom(M,A), c(A), - 22 <= #count{ B : aux_carbon_reachable(A,B) }. - -A property that needs computed facts (molecular weight, or rings larger than 8): - name: aux_has_macrocycle - description: contains a ring larger than 8 atoms - program: aux_has_macrocycle(M) :- ring_size(M,S), S > 8. - (weight variant: aux_high_molecular_weight(M) :- mol_weight(M,W), W >= 500.) - -Guidance (learned from failure analysis): -- Prefer small predicates identifying specific molecular features. The ILP system will later - combine them into larger conjunctions. Favor predicates that can be reused across many - classes. -- Counting and absence are the highest-value predicates: a close sibling often differs only in - the NUMBER of a group (N carbons, N sugar units, one vs two carboxyls). Reach for #count / not. -- A carboxyl group's -OH oxygen also carries one hydrogen: a "hydroxy oxygen" test written as - o(O), has_1_hs(O) matches it too. Exclude the carboxyl carbon when counting hydroxy groups. -- Rings of size <=8 use in_ringN/ringN; rings LARGER than 8 must use ring_size(M,S), S>8. -- Prefer predicates TRUE for many molecules of the target class and FALSE for other molecules.\ +- `not p(X)` is valid, `not (p(X), q(X))` is NOT. + To negate a conjunction, define a helper predicate. + +There is a SHARED LIBRARY of auxiliary rules already written for other classes. REUSE them where possible. """ # Header comments the model may repeat inside "program"; the pipeline synthesizes them. _HEADER_RE = re.compile(r"^\s*%\s*(PREDICATE_NAME|DESCRIPTION)\s*:", re.IGNORECASE) +# Validate-and-reject grounds at most ``validate_max`` positives and negatives, which takes +# well under a second; anything near this ceiling is pathological rather than merely slow. +_VALIDATION_GROUNDING_TIMEOUT = 60.0 + +# Rejection codes for the pipeline's own gates, alongside the static-analysis codes from +# auxiliary_rules. Every rejection carries one, so a repair pass can branch on the failure +# type rather than reading the message. +ERROR_UNPARSEABLE = "unparseable" +ERROR_CLINGO_SYNTAX = "clingo_syntax" +ERROR_NO_GROUNDING = "no_grounding" +ERROR_DEGENERATE = "degenerate" +ERROR_NO_VALIDATION_MOLECULES = "no_validation_molecules" + def _load_train_samples(chebi_id, problem_dir, molecules, max_pos, max_neg): """Return ``(pos_rows, neg_rows)`` DataFrames from the class's train exs.pl.""" @@ -309,8 +250,7 @@ def build_user_prompt(self, chebi_id, info, candidates, ctx) -> str: {format_candidates(candidates)} Choose up to {self.n_predicates} auxiliary predicates that would help distinguish "{info['name']}" from other molecules. REUSE the candidates above wherever they -fit; only write NEW rules for properties they do not already cover. Favour count/absence and -recursion predicates, which the ILP system cannot express on its own. +fit; only write NEW rules for properties they do not already cover. {OUTPUT_CONTRACT} "program" is the clause text only — one molecule-level head plus any helper clauses. @@ -331,50 +271,107 @@ def prepare(self, blocks, ctx) -> None: the syntax gate simply derives nothing and is rejected as degenerate below. """ ctx["errors"] = {} + ctx["error_codes"] = {} ctx["progs"] = {} for label, source in blocks: prog = parse_rule_program(source, source_file=label) if prog is None: ctx["errors"][label] = "unparseable program" + ctx["error_codes"][label] = ERROR_UNPARSEABLE continue error = rule_program_error(prog.source) if error is not None: ctx["errors"][label] = f"clingo error: {error}" + ctx["error_codes"][label] = ERROR_CLINGO_SYNTAX continue ctx["progs"][label] = prog - context = [p for p in (_load_library_rule(s, self.library_dir) for s in ctx["reused_stems"]) if p] + # Keyed by stem so accept_reused can find each reuse's own extension afterwards. + reused_progs = {} + for stem in ctx["reused_stems"]: + prog = _load_library_rule(stem, self.library_dir) + if prog is not None: + reused_progs[stem] = prog + ctx["reused_progs"] = reused_progs + ctx["extensions_valid"] = False + context = list(reused_progs.values()) + + # Static gates that must run before clingo sees the programs, because grounding either + # cannot survive the failure (non-terminating recursion runs the machine out of memory) + # or cannot detect it (a cross product grounds "successfully", just wrongly and + # enormously). Judged over the whole set, since a recursion cycle may close through a + # sibling or a reused program. + static = static_rule_errors(context + list(ctx["progs"].values())) + # A reused program blamed for non-terminating recursion means the NEW programs are what + # closed the cycle — the library grounds without them — so they go as a group. A cross + # product is local to one clause, so a reused program carrying one says nothing about + # the new programs and must not take them down with it. + via_reused = sorted( + name for name in {p.name for p in context} & set(static) + if static[name][0] == ERROR_UNBOUNDED_RECURSION + ) + for label, prog in list(ctx["progs"].items()): + entry = static.get(prog.name) + if entry is None and via_reused: + entry = (ERROR_UNBOUNDED_RECURSION, + f"non-terminating recursion through reused {via_reused[0]}") + if entry is not None: + ctx["error_codes"][label], ctx["errors"][label] = entry + del ctx["progs"][label] + together = context + list(ctx["progs"].values()) if not together or not ctx["val_ids"]: ctx["extensions"] = {} return try: - ctx["extensions"] = derive_rule_extensions(together, ctx["val_facts"], ctx["val_ids"]) + ctx["extensions"] = derive_rule_extensions( + together, ctx["val_facts"], ctx["val_ids"], timeout=_VALIDATION_GROUNDING_TIMEOUT + ) + ctx["extensions_valid"] = True except Exception as e: # The set as a whole will not ground (unstratified negation across programs, say). # Nothing can be attributed, so every program is rejected with the shared cause. ctx["extensions"] = {} for label in ctx["progs"]: ctx["errors"][label] = f"class rules do not ground together: {str(e).strip().splitlines()[0]}" + ctx["error_codes"][label] = ERROR_NO_GROUNDING ctx["progs"] = {} - def accept(self, source, label, ctx) -> tuple[bool, str]: - if label in ctx["errors"]: - return False, ctx["errors"][label] - prog = ctx["progs"][label] + def _extension_gate(self, prog, key, ctx) -> tuple[bool, str]: + """Judge one program by the extension it derived over the class's train molecules.""" if not ctx["val_ids"]: + ctx["error_codes"][key] = ERROR_NO_VALIDATION_MOLECULES return False, "no validation molecules" by_mol = ctx["extensions"].get(prog.name, {}) n = len(ctx["val_ids"]) frac = len(by_mol) / n if frac == 0.0: + ctx["error_codes"][key] = ERROR_DEGENERATE return False, f"degenerate: fires on 0% of {n} train molecules" # A molecule-level flag true of every molecule carries no information. An atom- or # pair-level predicate that fires everywhere still says WHICH atoms, so it is kept. if frac >= 0.95 and _is_molecule_level(by_mol): + ctx["error_codes"][key] = ERROR_DEGENERATE return False, f"degenerate: molecule-level flag on {frac:.0%} of {n} train molecules" return True, f"fires on {frac:.0%}" + def accept(self, source, label, ctx) -> tuple[bool, str]: + if label in ctx["errors"]: + return False, ctx["errors"][label] + return self._extension_gate(ctx["progs"][label], label, ctx) + + def accept_reused(self, stem, ctx) -> tuple[bool, str]: + prog = ctx.get("reused_progs", {}).get(stem) + # No grounding means no evidence about the reuse — the new programs already carry the + # shared failure, and dropping the reuses on top of it would report a cause that was + # never measured. + if prog is None or not ctx.get("extensions_valid"): + return True, "" + return self._extension_gate(prog, stem, ctx) + + def rejection_code(self, label, ctx) -> str | None: + return ctx.get("error_codes", {}).get(label) + def add_to_library(self, source, chebi_id): return add_rule_to_library(source, library_dir=self.library_dir, chebi_id=chebi_id) diff --git a/chebILP/predicate_generation/llm_client.py b/chebILP/predicate_generation/llm_client.py index 5a48153..1b3c865 100644 --- a/chebILP/predicate_generation/llm_client.py +++ b/chebILP/predicate_generation/llm_client.py @@ -20,16 +20,22 @@ litellm.enable_json_schema_validation = True +# Anthropic models LiteLLM's own allowlist misses. Substrings are version-specific +# ("opus-5" does not match "opus-4-5"), so an older model never lands here by accident. +_NATIVE_STRUCTURED_OUTPUT_FAMILIES = ("fable", "mythos", "haiku", "opus-5", "sonnet-5") + + def _patch_litellm_native_structured_output() -> None: """Route always-on-thinking Anthropic models through native structured output. LiteLLM chooses the Anthropic structured-output path from a hardcoded model - allowlist that predates Fable 5 / Mythos 5, so those fall back to tool-based - JSON coercion with a forced ``tool_choice``. That is incompatible with their - always-on thinking (the API can't force a tool while thinking is active), so - the model returns no tool call and the response parses to an empty ``{}``. - Reroute them to the native ``output_config.format`` path LiteLLM already uses - for 4.6/4.7, which is thinking-compatible. + allowlist that stops at Opus 4.7, so anything newer falls back to tool-based + JSON coercion with a forced ``tool_choice``. That is incompatible with always-on + thinking (the API can't force a tool while thinking is active), and the reply + comes back malformed: an empty ``{}`` on Fable/Mythos, or the whole answer + stuffed as a string under the literal placeholder key ``$PARAMETER_NAME`` on + Opus 5. Reroute them to the native ``output_config.format`` path LiteLLM already + uses for 4.6/4.7, which is thinking-compatible. """ from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -44,7 +50,7 @@ def patched(self, non_default_params, optional_params, model, drop_params): if ( isinstance(response_format, dict) and "output_format" not in params - and any(family in model for family in ("fable", "mythos")) + and any(family in model for family in _NATIVE_STRUCTURED_OUTPUT_FAMILIES) ): output_format = self.map_response_format_to_anthropic_output_format(response_format) if output_format is not None: @@ -93,14 +99,20 @@ def structured_completion( max_retries: int = 5, api_base: str | None = None, ): - """Ask ``model`` for one structured answer. Returns ``(parsed, raw_json_text)``. - - ``raw`` is the model's JSON string, kept for the exchange log. Retries transient - connection/timeout/rate-limit errors with exponential backoff, and reasks when a - (typically weaker) model returns malformed or schema-invalid JSON. + """Ask ``model`` for one structured answer. Returns ``(parsed, raw_json_text, attempts)``. + + ``raw`` is the model's JSON string, kept for the exchange log. ``attempts`` is one + record per LLM call made (each ``{"error", "raw", "cost"}``), in order — the final + entry is the successful call (``error`` is ``None``); any earlier entries are reasks. + Retries transient connection/timeout/rate-limit errors with exponential backoff, and + reasks when a (typically weaker) model returns malformed or schema-invalid JSON. On + total failure the collected attempts are attached to the raised exception as + ``_chebilp_attempts`` so the caller can still log them. """ last_exc = None + attempts: list[dict] = [] for attempt in range(max_retries): + cost = None try: response = litellm.completion( model=model, @@ -113,6 +125,7 @@ def structured_completion( response_format=schema, api_base=api_base, ) + cost = getattr(response, "_hidden_params", {}).get("response_cost") choice = response.choices[0] raw = choice.message.content if choice.finish_reason == "content_filter": @@ -124,7 +137,9 @@ def structured_completion( f"empty completion content (finish_reason={choice.finish_reason}, " f"reasoning_chars={len(reasoning)}); raw Anthropic blocks: {original!r}" ) - return schema.model_validate_json(raw), raw + parsed = schema.model_validate_json(raw) + attempts.append({"error": None, "raw": raw, "cost": cost}) + return parsed, raw, attempts except _TRANSIENT as e: last_exc = e wait = 2 ** attempt @@ -133,6 +148,11 @@ def structured_completion( except (ValidationError, ValueError, litellm.JSONSchemaValidationError) as e: last_exc = e raw = getattr(e, "raw_response", None) + attempts.append({"error": str(e), "raw": raw, "cost": cost}) detail = f"\n raw response: {raw!r}" if raw else "" print(f" Invalid structured output (attempt {attempt + 1}/{max_retries}), reasking: {e}{detail}") + try: + last_exc._chebilp_attempts = attempts + except (AttributeError, TypeError): + pass # some exception types (e.g. pydantic's) forbid attribute assignment raise last_exc