sqllocks-shape-behavior 0.9.1__py3-none-any.whl
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.
- shape_behavior/__init__.py +40 -0
- shape_behavior/behaviors.py +150 -0
- shape_behavior/cli.py +240 -0
- shape_behavior/conditions.py +221 -0
- shape_behavior/dist.py +89 -0
- shape_behavior/events.py +108 -0
- shape_behavior/examples/equipment_maintenance.json +169 -0
- shape_behavior/examples/healthcare_screening.json +255 -0
- shape_behavior/examples/subscription.json +261 -0
- shape_behavior/extension.py +109 -0
- shape_behavior/gmf.py +417 -0
- shape_behavior/model.py +495 -0
- shape_behavior/params.py +85 -0
- shape_behavior/population.py +224 -0
- shape_behavior/primitives.py +734 -0
- shape_behavior/rng.py +47 -0
- shape_behavior/simulator.py +551 -0
- shape_behavior/timeutil.py +81 -0
- sqllocks_shape_behavior-0.9.1.dist-info/METADATA +48 -0
- sqllocks_shape_behavior-0.9.1.dist-info/RECORD +24 -0
- sqllocks_shape_behavior-0.9.1.dist-info/WHEEL +5 -0
- sqllocks_shape_behavior-0.9.1.dist-info/entry_points.txt +12 -0
- sqllocks_shape_behavior-0.9.1.dist-info/licenses/LICENSE +21 -0
- sqllocks_shape_behavior-0.9.1.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
"""Shape plugin: behavior models (state machines on a virtual clock).
|
|
2
|
+
|
|
3
|
+
The public names are importable from here, lazily: importing the package loads no module of
|
|
4
|
+
this plugin beyond this file, and none of NumPy or Arrow. See ``docs/plugins/behavior.md``.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import importlib
|
|
10
|
+
from typing import Any
|
|
11
|
+
|
|
12
|
+
SHAPE_API = "1.0"
|
|
13
|
+
|
|
14
|
+
_EXPORTS = {
|
|
15
|
+
"Module": "shape_behavior.model",
|
|
16
|
+
"ModuleError": "shape_behavior.model",
|
|
17
|
+
"load_module": "shape_behavior.model",
|
|
18
|
+
"Population": "shape_behavior.population",
|
|
19
|
+
"SimConfig": "shape_behavior.simulator",
|
|
20
|
+
"Simulator": "shape_behavior.simulator",
|
|
21
|
+
"Checkpoint": "shape_behavior.simulator",
|
|
22
|
+
"EVENT_SCHEMA": "shape_behavior.events",
|
|
23
|
+
"import_gmf": "shape_behavior.gmf",
|
|
24
|
+
"ImportResult": "shape_behavior.gmf",
|
|
25
|
+
"Unsupported": "shape_behavior.gmf",
|
|
26
|
+
"UnsupportedGmfError": "shape_behavior.gmf",
|
|
27
|
+
"register_state_type": "shape_behavior.extension",
|
|
28
|
+
"StateHandler": "shape_behavior.extension",
|
|
29
|
+
"StateContext": "shape_behavior.extension",
|
|
30
|
+
"Emission": "shape_behavior.extension",
|
|
31
|
+
}
|
|
32
|
+
|
|
33
|
+
__all__ = sorted(_EXPORTS)
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def __getattr__(name: str) -> Any:
|
|
37
|
+
module = _EXPORTS.get(name)
|
|
38
|
+
if module is None:
|
|
39
|
+
raise AttributeError(f"module 'shape_behavior' has no attribute {name!r}")
|
|
40
|
+
return getattr(importlib.import_module(module), name)
|
|
@@ -0,0 +1,150 @@
|
|
|
1
|
+
"""Behavior modules as ``shape.behaviors`` plugins (plugin API v1, ``Behavior`` protocol).
|
|
2
|
+
|
|
3
|
+
A :class:`ModuleBehavior` wraps one :class:`~shape_behavior.model.Module` so the engine, the
|
|
4
|
+
plugin host, ``shape plugins list`` and the conformance kit can all reach it by name::
|
|
5
|
+
|
|
6
|
+
from shape_behavior.behaviors import behavior
|
|
7
|
+
|
|
8
|
+
b = behavior({"format": "shape-behavior/1", "name": "my_process", "states": {...}})
|
|
9
|
+
events = b.simulate(population=1000, seed=7, years=3) # a pyarrow.Table
|
|
10
|
+
|
|
11
|
+
The built-in examples register under ``shape.behaviors`` through the classes at the bottom of
|
|
12
|
+
this file; a third-party package registers its own module the same way
|
|
13
|
+
(``docs/plugins/behavior.md``, section 10).
|
|
14
|
+
"""
|
|
15
|
+
|
|
16
|
+
from __future__ import annotations
|
|
17
|
+
|
|
18
|
+
from typing import Any
|
|
19
|
+
|
|
20
|
+
from shape_behavior.extension import get_handler
|
|
21
|
+
from shape_behavior.model import EVENT_KINDS, Module, load_module
|
|
22
|
+
|
|
23
|
+
SHAPE_API = "1.0"
|
|
24
|
+
VERSION = "1.0"
|
|
25
|
+
|
|
26
|
+
DEFAULT_POPULATION: dict[str, Any] = {
|
|
27
|
+
"age_at_start": {"kind": "uniform", "low": 0, "high": 90},
|
|
28
|
+
"attributes": {"gender": {"kind": "categorical", "values": {"F": 0.5, "M": 0.5}}},
|
|
29
|
+
}
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def population_settings(modules: list[Module]) -> dict[str, Any]:
|
|
33
|
+
"""The population settings the modules ask for (``population_defaults``; the first module
|
|
34
|
+
wins on a clash), or :data:`DEFAULT_POPULATION` when none asks."""
|
|
35
|
+
import copy
|
|
36
|
+
|
|
37
|
+
spec: dict[str, Any] = {}
|
|
38
|
+
for m in reversed(modules):
|
|
39
|
+
spec.update(m.doc.get("population_defaults", {}))
|
|
40
|
+
return spec or copy.deepcopy(DEFAULT_POPULATION)
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
class ModuleBehavior:
|
|
44
|
+
"""One module as a ``shape.behaviors`` plugin: name, version, states, attributes, events."""
|
|
45
|
+
|
|
46
|
+
def __init__(self, module: Module | dict[str, Any] | str, version: str = VERSION) -> None:
|
|
47
|
+
self.module = load_module(module)
|
|
48
|
+
self.name: str = self.module.name
|
|
49
|
+
self.version: str = str(self.module.doc.get("version", version))
|
|
50
|
+
self.states: list[str] = list(self.module.states)
|
|
51
|
+
self.attributes: list[str] = sorted(self.module.attribute_uses())
|
|
52
|
+
self.events: list[str] = _event_kinds(self.module)
|
|
53
|
+
|
|
54
|
+
def simulate(self, population: int, seed: int, years: float) -> Any:
|
|
55
|
+
"""Events of ``population`` entities over ``years`` from 2020-01-01, deterministic per
|
|
56
|
+
``seed`` (an Arrow table in :data:`shape_behavior.EVENT_SCHEMA`)."""
|
|
57
|
+
from shape_behavior.population import Population
|
|
58
|
+
from shape_behavior.simulator import SimConfig, Simulator
|
|
59
|
+
from shape_behavior.timeutil import add_years, to_us
|
|
60
|
+
|
|
61
|
+
start = "2020-01-01"
|
|
62
|
+
pop = Population.from_dict(population_settings([self.module]), size=population, start=start)
|
|
63
|
+
sim = Simulator([self.module], pop, SimConfig(seed=seed))
|
|
64
|
+
return sim.run_until(add_years(to_us(start), years))
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def behavior(module: Module | dict[str, Any] | str, version: str = VERSION) -> ModuleBehavior:
|
|
68
|
+
"""A :class:`ModuleBehavior` for a module, a module document or a module file path."""
|
|
69
|
+
return ModuleBehavior(module, version)
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def _event_kinds(module: Module) -> list[str]:
|
|
73
|
+
kinds: set[str] = set()
|
|
74
|
+
for state in module.states.values():
|
|
75
|
+
kind = state["type"]
|
|
76
|
+
if kind == "event":
|
|
77
|
+
kinds.add(str(state["event"]))
|
|
78
|
+
elif kind in EVENT_KINDS:
|
|
79
|
+
kinds.add(EVENT_KINDS[kind])
|
|
80
|
+
else: # a registered state type may declare the kinds it emits (optional `kinds`)
|
|
81
|
+
declared = getattr(get_handler(kind), "kinds", None)
|
|
82
|
+
if declared is not None:
|
|
83
|
+
kinds.update(declared(state))
|
|
84
|
+
return sorted(kinds)
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
class Subscription(ModuleBehavior):
|
|
88
|
+
"""Subscription lifecycle: free trial, conversion, monthly payments, pause and cancel."""
|
|
89
|
+
|
|
90
|
+
def __init__(self) -> None:
|
|
91
|
+
super().__init__("subscription")
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
class EquipmentMaintenance(ModuleBehavior):
|
|
95
|
+
"""Equipment maintenance: run, degrade, fail, repair, with scheduled inspections."""
|
|
96
|
+
|
|
97
|
+
def __init__(self) -> None:
|
|
98
|
+
super().__init__("equipment_maintenance")
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
class HealthcareScreening(ModuleBehavior):
|
|
102
|
+
"""A tiny healthcare example (screening, diagnosis, treatment), written from scratch."""
|
|
103
|
+
|
|
104
|
+
def __init__(self) -> None:
|
|
105
|
+
super().__init__("healthcare_screening")
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
class EventSequence(ModuleBehavior):
|
|
109
|
+
"""An ordered funnel of named events with a continue probability per step (defaults)."""
|
|
110
|
+
|
|
111
|
+
def __init__(self) -> None:
|
|
112
|
+
from shape_behavior.primitives import event_sequence
|
|
113
|
+
|
|
114
|
+
super().__init__(event_sequence())
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
class TelemetrySeries(ModuleBehavior):
|
|
118
|
+
"""A regular reading per device with drift, noise, skipped and stuck readings (defaults)."""
|
|
119
|
+
|
|
120
|
+
def __init__(self) -> None:
|
|
121
|
+
from shape_behavior.primitives import telemetry_series
|
|
122
|
+
|
|
123
|
+
super().__init__(telemetry_series())
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
class TransactionStream(ModuleBehavior):
|
|
127
|
+
"""Transactions as a Poisson process per entity, with refunds and reversals (defaults)."""
|
|
128
|
+
|
|
129
|
+
def __init__(self) -> None:
|
|
130
|
+
from shape_behavior.primitives import transaction_stream
|
|
131
|
+
|
|
132
|
+
super().__init__(transaction_stream())
|
|
133
|
+
|
|
134
|
+
|
|
135
|
+
class FileArrival(ModuleBehavior):
|
|
136
|
+
"""A feed with an expected arrival per schedule slot: on time, late, missing, duplicated."""
|
|
137
|
+
|
|
138
|
+
def __init__(self) -> None:
|
|
139
|
+
from shape_behavior.primitives import file_arrival
|
|
140
|
+
|
|
141
|
+
super().__init__(file_arrival())
|
|
142
|
+
|
|
143
|
+
|
|
144
|
+
class EntityLifecycle(ModuleBehavior):
|
|
145
|
+
"""Create, update and delete events per entity (change-data-capture style)."""
|
|
146
|
+
|
|
147
|
+
def __init__(self) -> None:
|
|
148
|
+
from shape_behavior.primitives import entity_lifecycle
|
|
149
|
+
|
|
150
|
+
super().__init__(entity_lifecycle())
|
shape_behavior/cli.py
ADDED
|
@@ -0,0 +1,240 @@
|
|
|
1
|
+
"""``shape behave``: run, check and import behavior modules.
|
|
2
|
+
|
|
3
|
+
Reference: ``docs/plugins/behavior.md``, section 7.
|
|
4
|
+
|
|
5
|
+
::
|
|
6
|
+
|
|
7
|
+
shape behave run subscription --population 10000 --years 3 --seed 7 -o out/
|
|
8
|
+
shape behave check my_module.json
|
|
9
|
+
shape behave import-gmf downloaded_module.json -o converted.json
|
|
10
|
+
shape behave examples -o modules/
|
|
11
|
+
|
|
12
|
+
Nothing heavy loads at import time: NumPy, Arrow and the simulator load when a command runs.
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
from __future__ import annotations
|
|
16
|
+
|
|
17
|
+
import json
|
|
18
|
+
import sys
|
|
19
|
+
import time
|
|
20
|
+
from pathlib import Path
|
|
21
|
+
from typing import Any
|
|
22
|
+
|
|
23
|
+
from shape_behavior.params import (
|
|
24
|
+
PRIMITIVE_NAMES,
|
|
25
|
+
PrimitiveParamsError,
|
|
26
|
+
read_params,
|
|
27
|
+
unknown_primitive,
|
|
28
|
+
)
|
|
29
|
+
|
|
30
|
+
SHAPE_API = "1.0"
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
class BehaveCommand:
|
|
34
|
+
"""``shape behave COMMAND ...``."""
|
|
35
|
+
|
|
36
|
+
name = "behave"
|
|
37
|
+
help = "run state-machine behavior modules on a virtual clock (events as Parquet)"
|
|
38
|
+
|
|
39
|
+
def configure(self, parser: Any) -> None:
|
|
40
|
+
sub = parser.add_subparsers(dest="behave_cmd", required=True, metavar="COMMAND")
|
|
41
|
+
run = sub.add_parser("run", help="simulate a population and write the event stream")
|
|
42
|
+
run.add_argument(
|
|
43
|
+
"modules",
|
|
44
|
+
nargs="+",
|
|
45
|
+
metavar="MODULES",
|
|
46
|
+
help="module files (native or GMF JSON) or example names",
|
|
47
|
+
)
|
|
48
|
+
run.add_argument(
|
|
49
|
+
"--population", type=int, required=True, metavar="N", help="number of entities"
|
|
50
|
+
)
|
|
51
|
+
run.add_argument(
|
|
52
|
+
"--years", type=float, required=True, metavar="Y", help="years to simulate"
|
|
53
|
+
)
|
|
54
|
+
run.add_argument("--seed", type=int, default=0, help="random seed (default 0)")
|
|
55
|
+
run.add_argument("-o", "--out", required=True, metavar="OUT", help="output directory")
|
|
56
|
+
run.add_argument(
|
|
57
|
+
"--start", default="2020-01-01", help="virtual clock start (default 2020-01-01)"
|
|
58
|
+
)
|
|
59
|
+
run.add_argument(
|
|
60
|
+
"--population-spec",
|
|
61
|
+
metavar="FILE",
|
|
62
|
+
help="JSON population settings (attributes, age_at_start, arrival, lifetime)",
|
|
63
|
+
)
|
|
64
|
+
run.add_argument(
|
|
65
|
+
"--window-years", type=float, default=1.0, help="years per output file (default 1)"
|
|
66
|
+
)
|
|
67
|
+
run.add_argument(
|
|
68
|
+
"--poll", default="7 days", help="guard polling interval (default '7 days')"
|
|
69
|
+
)
|
|
70
|
+
run.add_argument("--strict", action="store_true", help="fail on unsupported GMF elements")
|
|
71
|
+
run.add_argument(
|
|
72
|
+
"--params",
|
|
73
|
+
metavar="FILE.json",
|
|
74
|
+
help="parameters of a primitive named in MODULES "
|
|
75
|
+
'({"format": "shape-behavior-params", "version": 1, "primitive": NAME, "params": {}})',
|
|
76
|
+
)
|
|
77
|
+
run.set_defaults(fn=_run)
|
|
78
|
+
check = sub.add_parser("check", help="validate modules; exit 1 when one is invalid")
|
|
79
|
+
check.add_argument("modules", nargs="+", metavar="MODULES")
|
|
80
|
+
check.add_argument("--strict", action="store_true", help="fail on unsupported GMF elements")
|
|
81
|
+
check.set_defaults(fn=_check)
|
|
82
|
+
imp = sub.add_parser(
|
|
83
|
+
"import-gmf", help="convert a GMF module and report what is unsupported"
|
|
84
|
+
)
|
|
85
|
+
imp.add_argument("file", help="the GMF JSON file you downloaded")
|
|
86
|
+
imp.add_argument("-o", "--out", metavar="OUT.json", help="write the native module here")
|
|
87
|
+
imp.add_argument("--strict", action="store_true", help="fail on unsupported elements")
|
|
88
|
+
imp.set_defaults(fn=_import)
|
|
89
|
+
ex = sub.add_parser("examples", help="list the built-in example modules, or write them")
|
|
90
|
+
ex.add_argument("-o", "--out", metavar="DIR", help="write the examples into DIR")
|
|
91
|
+
ex.set_defaults(fn=_examples)
|
|
92
|
+
|
|
93
|
+
def run(self, args: Any) -> int:
|
|
94
|
+
try:
|
|
95
|
+
return int(args.fn(args))
|
|
96
|
+
except (FileNotFoundError, json.JSONDecodeError, PrimitiveParamsError) as exc:
|
|
97
|
+
print(f"shape behave: {exc}", file=sys.stderr)
|
|
98
|
+
return 2
|
|
99
|
+
except ValueError as exc: # includes ModuleError, UnsupportedGmfError
|
|
100
|
+
print(f"shape behave: {exc}", file=sys.stderr)
|
|
101
|
+
return 1
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
def _load(sources: list[str], strict: bool, params: str | None = None) -> list[Any]:
|
|
105
|
+
from shape_behavior.model import load_module
|
|
106
|
+
|
|
107
|
+
primitive, values = read_params(Path(params)) if params else ("", {})
|
|
108
|
+
if params and primitive not in sources:
|
|
109
|
+
if primitive not in PRIMITIVE_NAMES:
|
|
110
|
+
raise unknown_primitive(primitive)
|
|
111
|
+
raise PrimitiveParamsError(
|
|
112
|
+
f"--params is for primitive {primitive!r}, which is not among the modules to run "
|
|
113
|
+
f"({', '.join(sources)})",
|
|
114
|
+
"primitive",
|
|
115
|
+
)
|
|
116
|
+
if params:
|
|
117
|
+
from shape_behavior.primitives import build
|
|
118
|
+
modules = []
|
|
119
|
+
for src in sources:
|
|
120
|
+
module = (
|
|
121
|
+
build(primitive, values)
|
|
122
|
+
if params and src == primitive
|
|
123
|
+
else load_module(src, strict=strict)
|
|
124
|
+
)
|
|
125
|
+
report = module.import_report
|
|
126
|
+
if report is not None and (report.unsupported or report.warnings):
|
|
127
|
+
print(f"{src}:\n{report.report()}", file=sys.stderr)
|
|
128
|
+
modules.append(module)
|
|
129
|
+
return modules
|
|
130
|
+
|
|
131
|
+
|
|
132
|
+
def _run(args: Any) -> int:
|
|
133
|
+
import pyarrow.parquet as pq # type: ignore[import-untyped]
|
|
134
|
+
|
|
135
|
+
from shape_behavior.behaviors import population_settings
|
|
136
|
+
from shape_behavior.population import Population
|
|
137
|
+
from shape_behavior.simulator import SimConfig, Simulator
|
|
138
|
+
from shape_behavior.timeutil import add_years, to_us
|
|
139
|
+
|
|
140
|
+
if args.population < 0 or args.years <= 0 or args.window_years <= 0:
|
|
141
|
+
print(
|
|
142
|
+
"shape behave: --population must be >= 0 and --years and --window-years > 0",
|
|
143
|
+
file=sys.stderr,
|
|
144
|
+
)
|
|
145
|
+
return 2
|
|
146
|
+
modules = _load(args.modules, args.strict, args.params)
|
|
147
|
+
if args.population_spec:
|
|
148
|
+
spec = json.loads(Path(args.population_spec).read_text(encoding="utf-8"))
|
|
149
|
+
else:
|
|
150
|
+
spec = population_settings(modules)
|
|
151
|
+
population = Population.from_dict(spec, size=args.population, start=args.start)
|
|
152
|
+
sim = Simulator(modules, population, SimConfig(seed=args.seed, poll=args.poll))
|
|
153
|
+
out = Path(args.out)
|
|
154
|
+
(out / "events").mkdir(parents=True, exist_ok=True)
|
|
155
|
+
start = to_us(args.start)
|
|
156
|
+
end = add_years(start, args.years)
|
|
157
|
+
started = time.perf_counter()
|
|
158
|
+
windows: list[dict[str, Any]] = []
|
|
159
|
+
cursor, part = start, 0
|
|
160
|
+
while cursor < end:
|
|
161
|
+
nxt = min(add_years(start, args.window_years * (part + 1)), end)
|
|
162
|
+
events = sim.run_until(nxt)
|
|
163
|
+
pq.write_table(events, out / "events" / f"part-{part:04d}.parquet", compression="snappy")
|
|
164
|
+
windows.append(
|
|
165
|
+
{
|
|
166
|
+
"file": f"events/part-{part:04d}.parquet",
|
|
167
|
+
"until": str(sim.now),
|
|
168
|
+
"rows": events.num_rows,
|
|
169
|
+
}
|
|
170
|
+
)
|
|
171
|
+
cursor, part = nxt, part + 1
|
|
172
|
+
pq.write_table(sim.entities(), out / "entities.parquet", compression="snappy")
|
|
173
|
+
elapsed = time.perf_counter() - started
|
|
174
|
+
total = sum(w["rows"] for w in windows)
|
|
175
|
+
manifest = {
|
|
176
|
+
"seed": args.seed,
|
|
177
|
+
"years": args.years,
|
|
178
|
+
"start": args.start,
|
|
179
|
+
"population": population.to_dict(),
|
|
180
|
+
"modules": [_module_entry(m, s) for m, s in zip(modules, args.modules, strict=True)],
|
|
181
|
+
"events": total,
|
|
182
|
+
"windows": windows,
|
|
183
|
+
"seconds": round(elapsed, 3),
|
|
184
|
+
"unsupported": [
|
|
185
|
+
str(u) for m in modules if m.import_report for u in m.import_report.unsupported
|
|
186
|
+
],
|
|
187
|
+
}
|
|
188
|
+
(out / "run.json").write_text(
|
|
189
|
+
json.dumps(manifest, indent=2) + "\n", encoding="utf-8", newline="\n"
|
|
190
|
+
)
|
|
191
|
+
print(f"{total:,} events for {args.population:,} entities in {elapsed:.1f}s -> {out}")
|
|
192
|
+
return 0
|
|
193
|
+
|
|
194
|
+
|
|
195
|
+
def _module_entry(module: Any, source: str) -> dict[str, Any]:
|
|
196
|
+
entry: dict[str, Any] = {"name": module.name, "digest": module.digest(), "source": source}
|
|
197
|
+
if "primitive" in module.doc: # a primitive: record the parameters it ran with
|
|
198
|
+
entry["primitive"] = module.doc["primitive"]
|
|
199
|
+
entry["parameters"] = module.doc["parameters"]
|
|
200
|
+
return entry
|
|
201
|
+
|
|
202
|
+
|
|
203
|
+
def _check(args: Any) -> int:
|
|
204
|
+
modules = _load(args.modules, args.strict)
|
|
205
|
+
for m in modules:
|
|
206
|
+
print(f"ok {m.name} ({len(m.states)} states)")
|
|
207
|
+
return 0
|
|
208
|
+
|
|
209
|
+
|
|
210
|
+
def _import(args: Any) -> int:
|
|
211
|
+
from shape_behavior.gmf import import_gmf
|
|
212
|
+
|
|
213
|
+
result = import_gmf(Path(args.file), strict=args.strict)
|
|
214
|
+
print(result.report())
|
|
215
|
+
if args.out:
|
|
216
|
+
Path(args.out).write_text(
|
|
217
|
+
json.dumps(result.module.to_dict(), indent=2) + "\n", encoding="utf-8", newline="\n"
|
|
218
|
+
)
|
|
219
|
+
print(f"wrote {args.out}")
|
|
220
|
+
return 0
|
|
221
|
+
|
|
222
|
+
|
|
223
|
+
def _examples(args: Any) -> int:
|
|
224
|
+
import importlib.resources
|
|
225
|
+
|
|
226
|
+
from shape_behavior.model import examples
|
|
227
|
+
|
|
228
|
+
names = examples()
|
|
229
|
+
if not args.out:
|
|
230
|
+
print("\n".join(names))
|
|
231
|
+
return 0
|
|
232
|
+
out = Path(args.out)
|
|
233
|
+
out.mkdir(parents=True, exist_ok=True)
|
|
234
|
+
for name in names:
|
|
235
|
+
text = (
|
|
236
|
+
importlib.resources.files("shape_behavior") / "examples" / f"{name}.json"
|
|
237
|
+
).read_text(encoding="utf-8")
|
|
238
|
+
(out / f"{name}.json").write_text(text, encoding="utf-8", newline="\n")
|
|
239
|
+
print(f"wrote {len(names)} modules to {out}")
|
|
240
|
+
return 0
|
|
@@ -0,0 +1,221 @@
|
|
|
1
|
+
"""Conditions: validation, and compilation into vectorized tests over rows of entities."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from collections.abc import Callable
|
|
6
|
+
from typing import Any
|
|
7
|
+
|
|
8
|
+
import numpy as np
|
|
9
|
+
|
|
10
|
+
from shape_behavior.timeutil import to_us, unit_us
|
|
11
|
+
|
|
12
|
+
OPS = ("==", "!=", "<", "<=", ">", ">=")
|
|
13
|
+
NIL_OPS = ("is_nil", "is_not_nil")
|
|
14
|
+
LEAF_TYPES = ("true", "false", "attribute", "age", "date", "active_condition", "active_medication")
|
|
15
|
+
GROUP_TYPES = ("and", "or", "at_least", "at_most")
|
|
16
|
+
_CMP: dict[str, Callable[[Any, Any], Any]] = {
|
|
17
|
+
"==": np.equal,
|
|
18
|
+
"!=": np.not_equal,
|
|
19
|
+
"<": np.less,
|
|
20
|
+
"<=": np.less_equal,
|
|
21
|
+
">": np.greater,
|
|
22
|
+
">=": np.greater_equal,
|
|
23
|
+
}
|
|
24
|
+
_Test = Callable[[Any, Any, Any], Any]
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def _num(v: Any) -> float | None:
|
|
28
|
+
if isinstance(v, bool):
|
|
29
|
+
return 1.0 if v else 0.0
|
|
30
|
+
if isinstance(v, (int, float)):
|
|
31
|
+
return float(v)
|
|
32
|
+
return None
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def check(doc: Any, where: str) -> list[str]:
|
|
36
|
+
"""Problems with a condition document, empty when it is fine."""
|
|
37
|
+
if not isinstance(doc, dict):
|
|
38
|
+
return [f"{where}: a condition must be an object"]
|
|
39
|
+
kind = doc.get("type")
|
|
40
|
+
if kind in ("true", "false"):
|
|
41
|
+
return []
|
|
42
|
+
if kind == "attribute":
|
|
43
|
+
out = [] if isinstance(doc.get("attribute"), str) else [f"{where}: needs 'attribute'"]
|
|
44
|
+
op = doc.get("op")
|
|
45
|
+
if op in NIL_OPS:
|
|
46
|
+
return out
|
|
47
|
+
if op not in OPS:
|
|
48
|
+
return [*out, f"{where}: op must be one of {', '.join(OPS + NIL_OPS)}"]
|
|
49
|
+
if "value" not in doc:
|
|
50
|
+
out.append(f"{where}: needs 'value'")
|
|
51
|
+
elif isinstance(doc["value"], str) and op not in ("==", "!="):
|
|
52
|
+
out.append(f"{where}: text attributes only compare with == and !=")
|
|
53
|
+
elif _num(doc["value"]) is None and not isinstance(doc["value"], str):
|
|
54
|
+
out.append(f"{where}: value must be a number, a bool or text")
|
|
55
|
+
return out
|
|
56
|
+
if kind in ("age", "date"):
|
|
57
|
+
out = []
|
|
58
|
+
if doc.get("op") not in OPS:
|
|
59
|
+
out.append(f"{where}: op must be one of {', '.join(OPS)}")
|
|
60
|
+
if "value" not in doc:
|
|
61
|
+
return [*out, f"{where}: needs 'value'"]
|
|
62
|
+
try:
|
|
63
|
+
if kind == "age":
|
|
64
|
+
if _num(doc["value"]) is None:
|
|
65
|
+
raise ValueError("age value must be a number")
|
|
66
|
+
unit_us(doc.get("unit", "years"))
|
|
67
|
+
else:
|
|
68
|
+
to_us(doc["value"])
|
|
69
|
+
except (ValueError, TypeError) as exc:
|
|
70
|
+
out.append(f"{where}: {exc}")
|
|
71
|
+
return out
|
|
72
|
+
if kind in ("active_condition", "active_medication"):
|
|
73
|
+
codes = doc.get("codes")
|
|
74
|
+
if not isinstance(codes, list) or not codes or not all(code_of(c) for c in codes):
|
|
75
|
+
return [f"{where}: needs a non-empty list of 'codes'"]
|
|
76
|
+
return []
|
|
77
|
+
if kind in GROUP_TYPES:
|
|
78
|
+
subs = doc.get("conditions")
|
|
79
|
+
if not isinstance(subs, list) or not subs:
|
|
80
|
+
return [f"{where}: needs a non-empty list of 'conditions'"]
|
|
81
|
+
out = []
|
|
82
|
+
if kind in ("at_least", "at_most"):
|
|
83
|
+
key = "minimum" if kind == "at_least" else "maximum"
|
|
84
|
+
if isinstance(doc.get(key), bool) or not isinstance(doc.get(key), int):
|
|
85
|
+
out.append(f"{where}: needs an integer {key!r}")
|
|
86
|
+
for i, sub in enumerate(subs):
|
|
87
|
+
out += check(sub, f"{where}.conditions[{i}]")
|
|
88
|
+
return out
|
|
89
|
+
if kind == "not":
|
|
90
|
+
return check(doc.get("condition"), f"{where}.condition")
|
|
91
|
+
return [f"{where}: unknown condition type {kind!r}"]
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
def code_of(c: Any) -> str | None:
|
|
95
|
+
"""The code string of a ``codes`` entry (a dict with ``code``, or the string itself)."""
|
|
96
|
+
if isinstance(c, str) and c:
|
|
97
|
+
return c
|
|
98
|
+
if isinstance(c, dict) and isinstance(c.get("code"), str) and c["code"]:
|
|
99
|
+
return str(c["code"])
|
|
100
|
+
return None
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
def walk(doc: Any) -> list[dict[str, Any]]:
|
|
104
|
+
"""Every condition node in ``doc``, the root first."""
|
|
105
|
+
out = [doc]
|
|
106
|
+
if doc.get("type") in GROUP_TYPES:
|
|
107
|
+
for sub in doc["conditions"]:
|
|
108
|
+
out += walk(sub)
|
|
109
|
+
elif doc.get("type") == "not":
|
|
110
|
+
out += walk(doc["condition"])
|
|
111
|
+
return out
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
class Compiled:
|
|
115
|
+
"""A compiled condition: ``test(store, rows, time)`` gives a bool array over the rows;
|
|
116
|
+
``ready`` (``None`` unless the condition depends only on age) gives the earliest time at
|
|
117
|
+
which it holds, ``INF`` for never."""
|
|
118
|
+
|
|
119
|
+
def __init__(self, test: _Test, ready: Callable[[Any, Any], Any] | None = None) -> None:
|
|
120
|
+
self.test = test
|
|
121
|
+
self.ready = ready
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
def compile_condition(doc: dict[str, Any], bit: Callable[[str, str], int]) -> Compiled:
|
|
125
|
+
"""Compile ``doc``; ``bit(kind, code)`` gives the bit of a code in the ``condition`` or
|
|
126
|
+
``medication`` active set."""
|
|
127
|
+
kind = doc["type"]
|
|
128
|
+
if kind == "true":
|
|
129
|
+
return Compiled(lambda s, r, t: np.ones(len(r), bool))
|
|
130
|
+
if kind == "false":
|
|
131
|
+
return Compiled(lambda s, r, t: np.zeros(len(r), bool))
|
|
132
|
+
if kind == "attribute":
|
|
133
|
+
return Compiled(_attribute_test(doc))
|
|
134
|
+
if kind == "age":
|
|
135
|
+
return _age(doc)
|
|
136
|
+
if kind == "date":
|
|
137
|
+
value = to_us(doc["value"])
|
|
138
|
+
cmp = _CMP[doc["op"]]
|
|
139
|
+
return Compiled(lambda s, r, t: cmp(t, value))
|
|
140
|
+
if kind in ("active_condition", "active_medication"):
|
|
141
|
+
mask = 0
|
|
142
|
+
which = "condition" if kind == "active_condition" else "medication"
|
|
143
|
+
for c in doc["codes"]:
|
|
144
|
+
mask |= 1 << bit(which, str(code_of(c)))
|
|
145
|
+
field = "cond_mask" if which == "condition" else "med_mask"
|
|
146
|
+
m = np.uint64(mask)
|
|
147
|
+
return Compiled(lambda s, r, t: (getattr(s, field)[r] & m) != 0)
|
|
148
|
+
if kind == "not":
|
|
149
|
+
inner = compile_condition(doc["condition"], bit).test
|
|
150
|
+
return Compiled(lambda s, r, t: ~inner(s, r, t))
|
|
151
|
+
subs = [compile_condition(c, bit) for c in doc["conditions"]]
|
|
152
|
+
tests = [c.test for c in subs]
|
|
153
|
+
if kind == "and":
|
|
154
|
+
ready = _and_ready(subs)
|
|
155
|
+
return Compiled(lambda s, r, t: np.logical_and.reduce([f(s, r, t) for f in tests]), ready)
|
|
156
|
+
if kind == "or":
|
|
157
|
+
return Compiled(lambda s, r, t: np.logical_or.reduce([f(s, r, t) for f in tests]))
|
|
158
|
+
count = doc.get("minimum" if kind == "at_least" else "maximum", 0)
|
|
159
|
+
|
|
160
|
+
def counted(s: Any, r: Any, t: Any) -> Any:
|
|
161
|
+
n = np.sum([f(s, r, t) for f in tests], axis=0)
|
|
162
|
+
return n >= count if kind == "at_least" else n <= count
|
|
163
|
+
|
|
164
|
+
return Compiled(counted)
|
|
165
|
+
|
|
166
|
+
|
|
167
|
+
def _and_ready(subs: list[Compiled]) -> Callable[[Any, Any], Any] | None:
|
|
168
|
+
readies = [c.ready for c in subs]
|
|
169
|
+
if any(r is None for r in readies):
|
|
170
|
+
return None
|
|
171
|
+
fns: list[Callable[[Any, Any], Any]] = [r for r in readies if r is not None]
|
|
172
|
+
return lambda s, r: np.maximum.reduce([f(s, r) for f in fns])
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
def _attribute_test(doc: dict[str, Any]) -> _Test:
|
|
176
|
+
name, op = doc["attribute"], doc["op"]
|
|
177
|
+
if op == "is_nil":
|
|
178
|
+
return lambda s, r, t: s.is_nil(name, r)
|
|
179
|
+
if op == "is_not_nil":
|
|
180
|
+
return lambda s, r, t: ~s.is_nil(name, r)
|
|
181
|
+
value = doc["value"]
|
|
182
|
+
cmp = _CMP[op]
|
|
183
|
+
if isinstance(value, str):
|
|
184
|
+
|
|
185
|
+
def text_test(s: Any, r: Any, t: Any) -> Any:
|
|
186
|
+
if name in s.num: # a numeric attribute never equals text
|
|
187
|
+
return np.full(len(r), op == "!=")
|
|
188
|
+
code = s.vindex.get(name, {}).get(value, -2)
|
|
189
|
+
return cmp(s.cat[name][r], code) if name in s.cat else np.full(len(r), op == "!=")
|
|
190
|
+
|
|
191
|
+
return text_test
|
|
192
|
+
num = float(_num(value) or 0.0)
|
|
193
|
+
|
|
194
|
+
def num_test(s: Any, r: Any, t: Any) -> Any:
|
|
195
|
+
if name not in s.num:
|
|
196
|
+
return np.zeros(len(r), bool)
|
|
197
|
+
x = s.num[name][r]
|
|
198
|
+
res = cmp(x, num)
|
|
199
|
+
return res | np.isnan(x) if op == "!=" else res
|
|
200
|
+
|
|
201
|
+
return num_test
|
|
202
|
+
|
|
203
|
+
|
|
204
|
+
def _age(doc: dict[str, Any]) -> Compiled:
|
|
205
|
+
threshold = round(float(doc["value"]) * unit_us(doc.get("unit", "years")))
|
|
206
|
+
op = doc["op"]
|
|
207
|
+
cmp = _CMP[op]
|
|
208
|
+
|
|
209
|
+
def test(s: Any, r: Any, t: Any) -> Any:
|
|
210
|
+
return cmp(t - s.born[r], threshold)
|
|
211
|
+
|
|
212
|
+
ready: Callable[[Any, Any], Any] | None = None
|
|
213
|
+
if op in (">=", ">"):
|
|
214
|
+
# the first microsecond at which age >= value (or > value)
|
|
215
|
+
def ready(s: Any, r: Any) -> Any:
|
|
216
|
+
return s.born[r] + threshold + (1 if op == ">" else 0)
|
|
217
|
+
|
|
218
|
+
elif op in ("<", "<="):
|
|
219
|
+
# holds from the start and stops: handled by the caller as "now or never"
|
|
220
|
+
ready = None
|
|
221
|
+
return Compiled(test, ready)
|