neurosymbolic-arithmetic 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 outlawsujith
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,148 @@
1
+ Metadata-Version: 2.4
2
+ Name: neurosymbolic-arithmetic
3
+ Version: 0.1.0
4
+ Summary: Neurosymbolic AI from scratch in NumPy: a neural net learns digits through differentiable symbolic reasoning, with an explainable logic reasoner.
5
+ Author: outlawsujith
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/outlawsujith/neurosymbolic-arithmetic
8
+ Project-URL: Issues, https://github.com/outlawsujith/neurosymbolic-arithmetic/issues
9
+ Keywords: neurosymbolic,neuro-symbolic,logic,datalog,semantic-loss,numpy
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Science/Research
12
+ Classifier: Programming Language :: Python :: 3
13
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
14
+ Requires-Python: >=3.9
15
+ Description-Content-Type: text/markdown
16
+ License-File: LICENSE
17
+ Requires-Dist: numpy>=1.22
18
+ Dynamic: license-file
19
+
20
+ # Neurosymbolic Arithmetic
21
+
22
+ A small, dependency-light neurosymbolic AI system written from scratch in NumPy.
23
+ A neural network learns to **read digits** from noisy images while being told only
24
+ the **result of a symbolic program** run on those digits (for example, their sum).
25
+ No digit is ever labelled. A logic reasoner then draws conclusions from what the
26
+ network perceives and explains each one with a proof.
27
+
28
+ ```
29
+ image 1 ──► neural net ──► P(digit 1) ─┐
30
+ ├─► symbolic program ──► P(sum) ──► loss vs. observed sum
31
+ image 2 ──► neural net ──► P(digit 2) ─┘ (exact probabilistic inference,
32
+ (shared) gradients flow back through it)
33
+ ```
34
+
35
+ ## Why this is neurosymbolic
36
+
37
+ | Part | File | Role |
38
+ |---|---|---|
39
+ | Neural perception | [nesy/perception.py](nesy/perception.py) | MLP with hand-written backprop and Adam; maps an image to a distribution over digits |
40
+ | Differentiable reasoning | [nesy/semantic.py](nesy/semantic.py) | Sums the probability of every world `(a, b)` that the symbolic program maps to the observed answer; `-log P(answer)` is the training loss |
41
+ | Symbolic reasoning | [nesy/logic.py](nesy/logic.py) | Datalog-style forward chainer with unification, arithmetic built-ins and proof traces |
42
+ | Pipeline | [nesy/pipeline.py](nesy/pipeline.py) | Training loops, a pure-neural baseline, and perceive-then-reason |
43
+ | Data | [nesy/data.py](nesy/data.py) | Synthetic noisy, randomly shifted digit glyphs (nothing to download) |
44
+
45
+ The training signal is the approach used by DeepProbLog and the semantic loss,
46
+ written out explicitly for programs over two digits:
47
+
48
+ ```
49
+ P(y) = Σ p1[a] · p2[b] over all digits a, b with f(a, b) = y
50
+ loss = −log P(y_observed)
51
+ ```
52
+
53
+ ## Quick start
54
+
55
+ Requires Python 3.9+ and NumPy.
56
+
57
+ ```bash
58
+ pip install git+https://github.com/outlawsujith/neurosymbolic-arithmetic.git
59
+ nesy-train # learn digits from sums, compare with a neural baseline
60
+ nesy-demo # perceive, reason, explain
61
+ ```
62
+
63
+ To work on the code instead, clone the repo and install it in editable mode:
64
+
65
+ ```bash
66
+ git clone https://github.com/outlawsujith/neurosymbolic-arithmetic.git
67
+ cd neurosymbolic-arithmetic
68
+ pip install -e .
69
+ python -m unittest discover -s tests -v
70
+ ```
71
+
72
+ If the `nesy-train` command is not on your PATH (or is blocked on your machine), use
73
+ `python -m nesy.train` and `python -m nesy.demo` instead.
74
+
75
+ Other tasks: `nesy-train --task product` (also `difference`, `concat`).
76
+ Add your own in `TASKS` in [nesy/semantic.py](nesy/semantic.py) as any Python function of two digits.
77
+
78
+ ## Results
79
+
80
+ `nesy-train --sweep` (seed 0, 2000 unseen test pairs, both models given the
81
+ same data and the same number of gradient steps):
82
+
83
+ | Training pairs | Neurosymbolic (sum accuracy) | Neural baseline (sum accuracy) | Digit accuracy (neurosymbolic) |
84
+ |---:|---:|---:|---:|
85
+ | 100 | 8.9% | 8.3% | 14.9% |
86
+ | 250 | 17.8% | 8.6% | 35.0% |
87
+ | 500 | 23.1% | 9.7% | 38.5% |
88
+ | 1000 | 95.3% | 10.8% | 97.6% |
89
+ | 2000 | 97.2% | 14.4% | 98.6% |
90
+ | 4000 | 97.7% | 26.9% | 98.9% |
91
+
92
+ The baseline is an MLP that takes both images and predicts the sum directly, so
93
+ it has to rediscover addition from examples. The neurosymbolic model is given
94
+ addition as knowledge and only has to learn perception. On the `product` task
95
+ with 2000 pairs the gap is 97.1% vs 37.9%.
96
+
97
+ Limits worth knowing: with fewer than about 1000 pairs the answer-only signal is
98
+ too weak and the neurosymbolic model also fails; these are single-seed numbers on
99
+ synthetic data; and exact inference enumerates all 100 digit pairs, which does not
100
+ scale to long programs without approximation.
101
+
102
+ ## Demo output
103
+
104
+ ```
105
+ Neural perception: left = 4 (100% confident), right = 5 (100% confident) [truth: 4, 5]
106
+
107
+ Symbolic reasoning derived:
108
+ odd_sum
109
+ product(20)
110
+ right_greater
111
+ sum(9)
112
+
113
+ Why does `odd_sum` hold?
114
+ odd_sum [by rule: odd_sum :- sum(Z), mod(Z, 2, 1).]
115
+ sum(9) [by rule: sum(Z) :- left(X), right(Y), add(X, Y, Z).]
116
+ left(4) [perceived]
117
+ right(5) [perceived]
118
+ add(4, 5, 9) [arithmetic]
119
+ mod(9, 2, 1) [arithmetic]
120
+
121
+ Abduction: told only that the true sum is 9, the most probable digits are (4, 5) with posterior 100%
122
+ ```
123
+
124
+ ## Using the reasoner on its own
125
+
126
+ ```python
127
+ from nesy import KnowledgeBase
128
+
129
+ kb = KnowledgeBase([
130
+ "path(X, Y) :- edge(X, Y).",
131
+ "path(X, Z) :- edge(X, Y), path(Y, Z).",
132
+ ])
133
+ for fact in ("edge(a, b)", "edge(b, c)"):
134
+ kb.add_fact(fact)
135
+ kb.forward_chain()
136
+ print(kb.holds("path(a, c)")) # True
137
+ print("\n".join(kb.explain("path(a, c)")))
138
+ ```
139
+
140
+ ## References
141
+
142
+ - Manhaeve et al., *DeepProbLog: Neural Probabilistic Logic Programming*, NeurIPS 2018
143
+ - Xu et al., *A Semantic Loss Function for Deep Learning with Symbolic Knowledge*, ICML 2018
144
+ - Garcez & Lamb, *Neurosymbolic AI: The 3rd Wave*, 2020
145
+
146
+ ## License
147
+
148
+ MIT. See [LICENSE](LICENSE).
@@ -0,0 +1,129 @@
1
+ # Neurosymbolic Arithmetic
2
+
3
+ A small, dependency-light neurosymbolic AI system written from scratch in NumPy.
4
+ A neural network learns to **read digits** from noisy images while being told only
5
+ the **result of a symbolic program** run on those digits (for example, their sum).
6
+ No digit is ever labelled. A logic reasoner then draws conclusions from what the
7
+ network perceives and explains each one with a proof.
8
+
9
+ ```
10
+ image 1 ──► neural net ──► P(digit 1) ─┐
11
+ ├─► symbolic program ──► P(sum) ──► loss vs. observed sum
12
+ image 2 ──► neural net ──► P(digit 2) ─┘ (exact probabilistic inference,
13
+ (shared) gradients flow back through it)
14
+ ```
15
+
16
+ ## Why this is neurosymbolic
17
+
18
+ | Part | File | Role |
19
+ |---|---|---|
20
+ | Neural perception | [nesy/perception.py](nesy/perception.py) | MLP with hand-written backprop and Adam; maps an image to a distribution over digits |
21
+ | Differentiable reasoning | [nesy/semantic.py](nesy/semantic.py) | Sums the probability of every world `(a, b)` that the symbolic program maps to the observed answer; `-log P(answer)` is the training loss |
22
+ | Symbolic reasoning | [nesy/logic.py](nesy/logic.py) | Datalog-style forward chainer with unification, arithmetic built-ins and proof traces |
23
+ | Pipeline | [nesy/pipeline.py](nesy/pipeline.py) | Training loops, a pure-neural baseline, and perceive-then-reason |
24
+ | Data | [nesy/data.py](nesy/data.py) | Synthetic noisy, randomly shifted digit glyphs (nothing to download) |
25
+
26
+ The training signal is the approach used by DeepProbLog and the semantic loss,
27
+ written out explicitly for programs over two digits:
28
+
29
+ ```
30
+ P(y) = Σ p1[a] · p2[b] over all digits a, b with f(a, b) = y
31
+ loss = −log P(y_observed)
32
+ ```
33
+
34
+ ## Quick start
35
+
36
+ Requires Python 3.9+ and NumPy.
37
+
38
+ ```bash
39
+ pip install git+https://github.com/outlawsujith/neurosymbolic-arithmetic.git
40
+ nesy-train # learn digits from sums, compare with a neural baseline
41
+ nesy-demo # perceive, reason, explain
42
+ ```
43
+
44
+ To work on the code instead, clone the repo and install it in editable mode:
45
+
46
+ ```bash
47
+ git clone https://github.com/outlawsujith/neurosymbolic-arithmetic.git
48
+ cd neurosymbolic-arithmetic
49
+ pip install -e .
50
+ python -m unittest discover -s tests -v
51
+ ```
52
+
53
+ If the `nesy-train` command is not on your PATH (or is blocked on your machine), use
54
+ `python -m nesy.train` and `python -m nesy.demo` instead.
55
+
56
+ Other tasks: `nesy-train --task product` (also `difference`, `concat`).
57
+ Add your own in `TASKS` in [nesy/semantic.py](nesy/semantic.py) as any Python function of two digits.
58
+
59
+ ## Results
60
+
61
+ `nesy-train --sweep` (seed 0, 2000 unseen test pairs, both models given the
62
+ same data and the same number of gradient steps):
63
+
64
+ | Training pairs | Neurosymbolic (sum accuracy) | Neural baseline (sum accuracy) | Digit accuracy (neurosymbolic) |
65
+ |---:|---:|---:|---:|
66
+ | 100 | 8.9% | 8.3% | 14.9% |
67
+ | 250 | 17.8% | 8.6% | 35.0% |
68
+ | 500 | 23.1% | 9.7% | 38.5% |
69
+ | 1000 | 95.3% | 10.8% | 97.6% |
70
+ | 2000 | 97.2% | 14.4% | 98.6% |
71
+ | 4000 | 97.7% | 26.9% | 98.9% |
72
+
73
+ The baseline is an MLP that takes both images and predicts the sum directly, so
74
+ it has to rediscover addition from examples. The neurosymbolic model is given
75
+ addition as knowledge and only has to learn perception. On the `product` task
76
+ with 2000 pairs the gap is 97.1% vs 37.9%.
77
+
78
+ Limits worth knowing: with fewer than about 1000 pairs the answer-only signal is
79
+ too weak and the neurosymbolic model also fails; these are single-seed numbers on
80
+ synthetic data; and exact inference enumerates all 100 digit pairs, which does not
81
+ scale to long programs without approximation.
82
+
83
+ ## Demo output
84
+
85
+ ```
86
+ Neural perception: left = 4 (100% confident), right = 5 (100% confident) [truth: 4, 5]
87
+
88
+ Symbolic reasoning derived:
89
+ odd_sum
90
+ product(20)
91
+ right_greater
92
+ sum(9)
93
+
94
+ Why does `odd_sum` hold?
95
+ odd_sum [by rule: odd_sum :- sum(Z), mod(Z, 2, 1).]
96
+ sum(9) [by rule: sum(Z) :- left(X), right(Y), add(X, Y, Z).]
97
+ left(4) [perceived]
98
+ right(5) [perceived]
99
+ add(4, 5, 9) [arithmetic]
100
+ mod(9, 2, 1) [arithmetic]
101
+
102
+ Abduction: told only that the true sum is 9, the most probable digits are (4, 5) with posterior 100%
103
+ ```
104
+
105
+ ## Using the reasoner on its own
106
+
107
+ ```python
108
+ from nesy import KnowledgeBase
109
+
110
+ kb = KnowledgeBase([
111
+ "path(X, Y) :- edge(X, Y).",
112
+ "path(X, Z) :- edge(X, Y), path(Y, Z).",
113
+ ])
114
+ for fact in ("edge(a, b)", "edge(b, c)"):
115
+ kb.add_fact(fact)
116
+ kb.forward_chain()
117
+ print(kb.holds("path(a, c)")) # True
118
+ print("\n".join(kb.explain("path(a, c)")))
119
+ ```
120
+
121
+ ## References
122
+
123
+ - Manhaeve et al., *DeepProbLog: Neural Probabilistic Logic Programming*, NeurIPS 2018
124
+ - Xu et al., *A Semantic Loss Function for Deep Learning with Symbolic Knowledge*, ICML 2018
125
+ - Garcez & Lamb, *Neurosymbolic AI: The 3rd Wave*, 2020
126
+
127
+ ## License
128
+
129
+ MIT. See [LICENSE](LICENSE).
@@ -0,0 +1,21 @@
1
+ """Neurosymbolic arithmetic: neural perception + symbolic reasoning in pure NumPy."""
2
+
3
+ __version__ = "0.1.0"
4
+
5
+ from .data import make_pairs, render_digit, to_ascii
6
+ from .logic import KnowledgeBase, parse_atom, parse_rule
7
+ from .perception import MLP, softmax
8
+ from .semantic import SymbolicTask, TASKS
9
+
10
+ __all__ = [
11
+ "make_pairs",
12
+ "render_digit",
13
+ "to_ascii",
14
+ "KnowledgeBase",
15
+ "parse_atom",
16
+ "parse_rule",
17
+ "MLP",
18
+ "softmax",
19
+ "SymbolicTask",
20
+ "TASKS",
21
+ ]
@@ -0,0 +1,63 @@
1
+ """Synthetic noisy digit images.
2
+
3
+ Each digit is a 5x7 bitmap glyph dropped at a random offset into a 9x7 canvas,
4
+ then corrupted with pixel flips and Gaussian noise. No download required, and
5
+ hard enough that perception has to be learned rather than looked up.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import numpy as np
11
+
12
+ GLYPHS = {
13
+ 0: [" ### ", "# #", "# ##", "# # #", "## #", "# #", " ### "],
14
+ 1: [" # ", " ## ", " # ", " # ", " # ", " # ", " ### "],
15
+ 2: [" ### ", "# #", " #", " # ", " # ", " # ", "#####"],
16
+ 3: [" ### ", "# #", " #", " ## ", " #", "# #", " ### "],
17
+ 4: [" # ", " ## ", " # # ", "# # ", "#####", " # ", " # "],
18
+ 5: ["#####", "# ", "#### ", " #", " #", "# #", " ### "],
19
+ 6: [" ### ", "# ", "# ", "#### ", "# #", "# #", " ### "],
20
+ 7: ["#####", " #", " # ", " # ", " # ", " # ", " # "],
21
+ 8: [" ### ", "# #", "# #", " ### ", "# #", "# #", " ### "],
22
+ 9: [" ### ", "# #", "# #", " ####", " #", " #", " ### "],
23
+ }
24
+
25
+ GLYPH_H, GLYPH_W = 7, 5
26
+ HEIGHT, WIDTH = 9, 7
27
+ IMAGE_DIM = HEIGHT * WIDTH
28
+
29
+ _BITMAPS = np.array(
30
+ [[[c == "#" for c in row] for row in GLYPHS[d]] for d in range(10)], dtype=float
31
+ )
32
+
33
+
34
+ def render_digit(digit: int, rng: np.random.Generator, noise: float = 0.2,
35
+ flip: float = 0.02) -> np.ndarray:
36
+ """Return one noisy image of `digit` as a flat vector of length IMAGE_DIM."""
37
+ canvas = np.zeros((HEIGHT, WIDTH))
38
+ dy = rng.integers(0, HEIGHT - GLYPH_H + 1)
39
+ dx = rng.integers(0, WIDTH - GLYPH_W + 1)
40
+ canvas[dy:dy + GLYPH_H, dx:dx + GLYPH_W] = _BITMAPS[digit]
41
+ flips = rng.random(canvas.shape) < flip
42
+ canvas = np.where(flips, 1.0 - canvas, canvas)
43
+ canvas = canvas + rng.normal(0.0, noise, canvas.shape)
44
+ return canvas.ravel()
45
+
46
+
47
+ def make_pairs(n: int, rng: np.random.Generator, **kwargs):
48
+ """Sample `n` pairs of digit images.
49
+
50
+ Returns (X1, X2, d1, d2): two image arrays of shape (n, IMAGE_DIM) and the
51
+ hidden digit labels. Training only ever sees a function of (d1, d2).
52
+ """
53
+ d1 = rng.integers(0, 10, n)
54
+ d2 = rng.integers(0, 10, n)
55
+ X1 = np.stack([render_digit(d, rng, **kwargs) for d in d1])
56
+ X2 = np.stack([render_digit(d, rng, **kwargs) for d in d2])
57
+ return X1, X2, d1, d2
58
+
59
+
60
+ def to_ascii(image: np.ndarray, threshold: float = 0.5) -> list[str]:
61
+ """Render a flat image as lines of text for terminal display."""
62
+ grid = image.reshape(HEIGHT, WIDTH)
63
+ return ["".join("#" if v > threshold else "." for v in row) for row in grid]
@@ -0,0 +1,73 @@
1
+ """Perceive two noisy digit images, reason about them, and explain the answer.
2
+
3
+ nesy-demo # uses model.npz, training one first if missing
4
+ nesy-demo --seed 7 --examples 3
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import argparse
10
+ import os
11
+
12
+ import numpy as np
13
+
14
+ from nesy import MLP, TASKS, make_pairs, to_ascii
15
+ from nesy.logic import format_atom
16
+ from nesy.pipeline import perceive_and_reason, train_neurosymbolic
17
+
18
+
19
+ def load_or_train(path: str) -> MLP:
20
+ if os.path.exists(path):
21
+ return MLP.load(path)
22
+ print("No saved model found; training one from sum labels (a few seconds)...\n")
23
+ task = TASKS["sum"]
24
+ X1, X2, d1, d2 = make_pairs(2000, np.random.default_rng(0))
25
+ net = train_neurosymbolic(task, X1, X2, task.encode(d1, d2), epochs=60)
26
+ net.save(path)
27
+ return net
28
+
29
+
30
+ def main():
31
+ parser = argparse.ArgumentParser(description=__doc__,
32
+ formatter_class=argparse.RawDescriptionHelpFormatter)
33
+ parser.add_argument("--model", default="model.npz")
34
+ parser.add_argument("--seed", type=int, default=1)
35
+ parser.add_argument("--examples", type=int, default=2)
36
+ args = parser.parse_args()
37
+
38
+ net = load_or_train(args.model)
39
+ X1, X2, d1, d2 = make_pairs(args.examples, np.random.default_rng(args.seed))
40
+
41
+ for i in range(args.examples):
42
+ print("=" * 64)
43
+ print(f"Example {i + 1}: the network sees two noisy images")
44
+ for a, b in zip(to_ascii(X1[i]), to_ascii(X2[i])):
45
+ print(f" {a} {b}")
46
+
47
+ kb, derived, (c1, c2) = perceive_and_reason(net, X1[i], X2[i])
48
+ left = kb.query("left(X)")[0]["X"]
49
+ right = kb.query("right(X)")[0]["X"]
50
+ print(f"\nNeural perception: left = {left} ({c1:.0%} confident), "
51
+ f"right = {right} ({c2:.0%} confident) [truth: {d1[i]}, {d2[i]}]")
52
+
53
+ print("\nSymbolic reasoning derived:")
54
+ for fact in sorted(derived, key=format_atom):
55
+ print(f" {format_atom(fact)}")
56
+
57
+ total = kb.query("sum(Z)")[0]["Z"]
58
+ target = "carry" if kb.holds("carry") else "even_sum" if kb.holds("even_sum") else "odd_sum"
59
+ print(f"\nWhy does `{target}` hold?")
60
+ for line in kb.explain(target):
61
+ print(f" {line}")
62
+
63
+ # Abduction: reason backwards from a known answer to the likeliest inputs.
64
+ p1 = net.predict_proba(X1[i][None])[0]
65
+ p2 = net.predict_proba(X2[i][None])[0]
66
+ a, b, post = TASKS["sum"].abduce(p1, p2, int(d1[i] + d2[i]))
67
+ print(f"\nAbduction: told only that the true sum is {d1[i] + d2[i]}, the most "
68
+ f"probable digits are ({a}, {b}) with posterior {post:.0%}")
69
+ print(f"Predicted sum: {total} True sum: {d1[i] + d2[i]}\n")
70
+
71
+
72
+ if __name__ == "__main__":
73
+ main()
@@ -0,0 +1,206 @@
1
+ """The symbolic half: a small Datalog-style forward-chaining reasoner.
2
+
3
+ Rules are Horn clauses written in Prolog-like syntax:
4
+
5
+ sum(Z) :- left(X), right(Y), add(X, Y, Z).
6
+
7
+ Names starting with an uppercase letter are variables; integers and lowercase
8
+ names are constants. Arithmetic is provided by built-in predicates. Every
9
+ derived fact remembers the rule and premises that produced it, so any
10
+ conclusion can be explained with a proof tree.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ import re
16
+ from dataclasses import dataclass
17
+ from typing import Callable, Iterator
18
+
19
+ Atom = tuple # (predicate, arg1, arg2, ...)
20
+
21
+
22
+ def is_var(term) -> bool:
23
+ return isinstance(term, str) and (term[0].isupper() or term[0] == "_")
24
+
25
+
26
+ def _parse_term(text: str):
27
+ text = text.strip()
28
+ return int(text) if re.fullmatch(r"-?\d+", text) else text
29
+
30
+
31
+ def parse_atom(text: str) -> Atom:
32
+ m = re.fullmatch(r"\s*([a-z]\w*)\s*(?:\((.*)\))?\s*", text)
33
+ if not m:
34
+ raise ValueError(f"cannot parse atom: {text!r}")
35
+ args = [_parse_term(a) for a in m.group(2).split(",")] if m.group(2) else []
36
+ return (m.group(1), *args)
37
+
38
+
39
+ def _split_body(text: str) -> list[str]:
40
+ parts, depth, current = [], 0, ""
41
+ for ch in text:
42
+ if ch == "," and depth == 0:
43
+ parts.append(current)
44
+ current = ""
45
+ continue
46
+ depth += ch == "("
47
+ depth -= ch == ")"
48
+ current += ch
49
+ return parts + [current]
50
+
51
+
52
+ @dataclass(frozen=True)
53
+ class Rule:
54
+ head: Atom
55
+ body: tuple
56
+
57
+ def __str__(self) -> str:
58
+ return f"{format_atom(self.head)} :- {', '.join(map(format_atom, self.body))}."
59
+
60
+
61
+ def parse_rule(text: str) -> Rule:
62
+ text = text.strip().rstrip(".")
63
+ head, sep, body = text.partition(":-")
64
+ if not sep:
65
+ raise ValueError(f"rule needs a body: {text!r}")
66
+ return Rule(parse_atom(head), tuple(parse_atom(a) for a in _split_body(body)))
67
+
68
+
69
+ def format_atom(atom: Atom) -> str:
70
+ pred, *args = atom
71
+ return f"{pred}({', '.join(map(str, args))})" if args else pred
72
+
73
+
74
+ def _builtin_add(x, y, z):
75
+ if x is not None and y is not None:
76
+ yield (x, y, x + y)
77
+ elif x is not None and z is not None:
78
+ yield (x, z - x, z)
79
+ elif y is not None and z is not None:
80
+ yield (z - y, y, z)
81
+
82
+
83
+ def _builtin_mul(x, y, z):
84
+ if x is not None and y is not None:
85
+ yield (x, y, x * y)
86
+
87
+
88
+ def _builtin_mod(x, y, z):
89
+ if x is not None and y is not None and y != 0:
90
+ yield (x, y, x % y)
91
+
92
+
93
+ def _comparison(test: Callable[[int, int], bool]):
94
+ def builtin(x, y):
95
+ if x is not None and y is not None and test(x, y):
96
+ yield (x, y)
97
+ return builtin
98
+
99
+
100
+ BUILTINS: dict[str, Callable] = {
101
+ "add": _builtin_add,
102
+ "mul": _builtin_mul,
103
+ "mod": _builtin_mod,
104
+ "lt": _comparison(lambda x, y: x < y),
105
+ "gt": _comparison(lambda x, y: x > y),
106
+ "le": _comparison(lambda x, y: x <= y),
107
+ "ge": _comparison(lambda x, y: x >= y),
108
+ "neq": _comparison(lambda x, y: x != y),
109
+ }
110
+
111
+
112
+ def _unify(pattern: Atom, fact: Atom, theta: dict):
113
+ """Match a possibly non-ground atom against a ground fact."""
114
+ if len(pattern) != len(fact) or pattern[0] != fact[0]:
115
+ return None
116
+ theta = dict(theta)
117
+ for p, f in zip(pattern[1:], fact[1:]):
118
+ if is_var(p):
119
+ if theta.setdefault(p, f) != f:
120
+ return None
121
+ elif p != f:
122
+ return None
123
+ return theta
124
+
125
+
126
+ class KnowledgeBase:
127
+ def __init__(self, rules=(), builtins: dict | None = None):
128
+ self.rules: list[Rule] = []
129
+ self.facts: set[Atom] = set()
130
+ # fact -> (rule, premises) for derived facts, None for given facts
131
+ self.provenance: dict[Atom, tuple | None] = {}
132
+ self.builtins = dict(BUILTINS if builtins is None else builtins)
133
+ for rule in rules:
134
+ self.add_rule(rule)
135
+
136
+ def add_rule(self, rule) -> None:
137
+ self.rules.append(parse_rule(rule) if isinstance(rule, str) else rule)
138
+
139
+ def add_fact(self, fact) -> None:
140
+ atom = parse_atom(fact) if isinstance(fact, str) else tuple(fact)
141
+ if any(is_var(t) for t in atom[1:]):
142
+ raise ValueError(f"facts must be ground: {format_atom(atom)}")
143
+ if atom not in self.facts:
144
+ self.facts.add(atom)
145
+ self.provenance[atom] = None
146
+
147
+ def _solve(self, body: tuple, theta: dict, support: tuple) -> Iterator[tuple]:
148
+ if not body:
149
+ yield theta, support
150
+ return
151
+ first, rest = body[0], body[1:]
152
+ pred = first[0]
153
+ if pred in self.builtins:
154
+ args = [theta.get(t) if is_var(t) else t for t in first[1:]]
155
+ candidates = [(pred, *result) for result in self.builtins[pred](*args)]
156
+ else:
157
+ candidates = [f for f in self.facts if f[0] == pred]
158
+ for fact in candidates:
159
+ new_theta = _unify(first, fact, theta)
160
+ if new_theta is not None:
161
+ yield from self._solve(rest, new_theta, support + (fact,))
162
+
163
+ def forward_chain(self) -> set[Atom]:
164
+ """Apply every rule until no new fact appears. Returns the new facts."""
165
+ derived: set[Atom] = set()
166
+ changed = True
167
+ while changed:
168
+ changed = False
169
+ for rule in self.rules:
170
+ for theta, support in list(self._solve(rule.body, {}, ())):
171
+ head = (rule.head[0], *(theta.get(t, t) for t in rule.head[1:]))
172
+ if any(is_var(t) for t in head[1:]):
173
+ raise ValueError(f"unsafe rule, unbound head variable: {rule}")
174
+ if head not in self.facts:
175
+ self.facts.add(head)
176
+ self.provenance[head] = (rule, support)
177
+ derived.add(head)
178
+ changed = True
179
+ return derived
180
+
181
+ def holds(self, query) -> bool:
182
+ atom = parse_atom(query) if isinstance(query, str) else tuple(query)
183
+ return atom in self.facts
184
+
185
+ def query(self, pattern) -> list[dict]:
186
+ """All variable bindings under which `pattern` matches a known fact."""
187
+ atom = parse_atom(pattern) if isinstance(pattern, str) else tuple(pattern)
188
+ matches = (_unify(atom, f, {}) for f in self.facts)
189
+ return [m for m in matches if m is not None]
190
+
191
+ def explain(self, fact, indent: int = 0) -> list[str]:
192
+ """Proof tree for a fact, as indented lines of text."""
193
+ atom = parse_atom(fact) if isinstance(fact, str) else tuple(fact)
194
+ pad = " " * indent
195
+ if atom[0] in self.builtins:
196
+ return [f"{pad}{format_atom(atom)} [arithmetic]"]
197
+ if atom not in self.facts:
198
+ return [f"{pad}{format_atom(atom)} [not provable]"]
199
+ source = self.provenance[atom]
200
+ if source is None:
201
+ return [f"{pad}{format_atom(atom)} [perceived]"]
202
+ rule, premises = source
203
+ lines = [f"{pad}{format_atom(atom)} [by rule: {rule}]"]
204
+ for premise in premises:
205
+ lines += self.explain(premise, indent + 1)
206
+ return lines