amachine.am_fast

  1from __future__ import annotations
  2
  3import sys
  4from pathlib import Path
  5import json
  6from typing import Any
  7
  8import pyarrow as pa
  9import pyarrow.parquet as pq
 10import numpy as np
 11import numpy.typing as npt
 12
 13from automata.fa.dfa import DFA
 14
 15_build_dir = Path(__file__).resolve().parent / "build"
 16
 17if str(_build_dir) not in sys.path:
 18    sys.path.insert(0, str(_build_dir))
 19
 20from ._am_fast import ( 
 21    generate_cpp, 
 22    block_entropy_convergence_cpp, 
 23    strongly_connected_components_cpp,
 24    minify_dfa_cpp
 25)
 26
 27from .json_utils import save_json
 28
 29def block_entropy_convergence(
 30    h_mu: float,
 31    n_states: int,
 32    n_symbols : int,
 33    convergence_tol: float,
 34    precision: float,
 35    eps: float,
 36    branches: list[tuple[float, list[float]]],
 37    trans: list[list[tuple[int, float, int]]],
 38    max_branches: int = 30_000_000
 39) -> Any :
 40    return block_entropy_convergence_cpp(
 41        h_mu            = float( h_mu ),
 42        n_states        = n_states,
 43        n_symbols       = n_symbols,
 44        convergence_tol = convergence_tol,
 45        precision       = float(precision),
 46        eps             = eps,
 47        branches        = branches,
 48        trans           = trans,
 49        max_branches    = max_branches
 50    )
 51
 52def strongly_connected_components( T ) :
 53    return strongly_connected_components_cpp( T )
 54
 55def generate_data(
 56    n_gen          : int,
 57    start_state    : int,
 58    transitions    : list[list[tuple[int, float, int]]],
 59    alphabet       : list[str],
 60    include_states : bool = False,
 61    random_seed    : int  = 42,
 62) -> dict[str, Any ]:
 63
 64    n_states  = len(transitions)
 65    n_symbols = len(alphabet)
 66
 67    if n_states > 65536:
 68        raise ValueError(f"Max states supported is 65536, got {n_states}")
 69    if n_symbols > 65536:
 70        raise ValueError(f"Max alphabet size supported is 65536, got {n_symbols}")
 71
 72    symbol_indices, state_indices = generate_cpp(
 73        n_gen             = n_gen,
 74        start_state_index = start_state,
 75        transitions       = transitions,
 76        include_states    = include_states,
 77        random_seed       = random_seed,
 78    )
 79
 80    res: dict[str, Any ] = {
 81        "symbol_index": np.from_dlpack(symbol_indices)
 82    }
 83
 84    if include_states:
 85        res["state_index"] = np.from_dlpack(state_indices)
 86
 87    return res
 88
 89def save_data(
 90    data             : dict[str, Any],
 91    file_prefix      : str,
 92    alphabet         : list[str],
 93    n_states         : int,
 94    start_state      : int,
 95    random_seed      : int,
 96    row_size         : int | None = None,
 97    machine_metadata : dict[str, Any] | None = None,
 98) -> None:
 99
100    def _uint_type(n: int) -> pa.DataType:
101        return pa.uint8() if n <= 256 else pa.uint16()
102
103    sym_type   = _uint_type(len(alphabet))
104    state_type = _uint_type(n_states)
105
106    columns: dict[str, pa.Array] = {
107        "symbol_index": pa.array(data["symbol_index"]).cast(sym_type)
108    }
109
110    final_state=start_state
111
112    if "state_index" in data:
113        state_data  = data["state_index"]
114        final_state = int(state_data[-1])
115        columns["state_index"] = pa.array(state_data[:-1]).cast(state_type)
116
117    iso_shifts = data.get("isomorphic_shifts", {})
118    iso_final_states: dict[int, int] = {}
119
120    for shift, shifted in iso_shifts.items():
121        columns[f"symbol_index_isoshift_{shift}"] = pa.array(shifted["symbol_index"]).cast(sym_type)
122        if "state_index" in shifted:
123            iso_state_data = shifted["state_index"]
124            iso_final_states[shift] = int(iso_state_data[-1])
125            columns[f"state_index_isoshift_{shift}"]= pa.array(iso_state_data[:-1]).cast(state_type)
126
127    parquet_meta: dict[str, Any] = {
128        "alphabet"          : alphabet,
129        "machine_metadata"  : machine_metadata or {},
130        "start_state"       : start_state,
131        "random_seed"       : random_seed,
132        "isomorphic_shifts" : sorted(iso_shifts.keys()),
133    }
134
135    if "state_index" in data:
136        parquet_meta["final_state"] = final_state
137
138    if iso_final_states:
139        parquet_meta["isomorphic_final_states"] = iso_final_states
140
141    if "belief_states" in data:
142        
143        belief_data  = data["belief_states"]
144        belief_dim   = belief_data.shape[1]
145
146        columns["belief_state"] = pa.FixedSizeListArray.from_arrays(
147            pa.array(belief_data[:-1].ravel(), type=pa.float32()),
148            belief_dim,
149        )
150
151        parquet_meta["belief_state_dim"]   = belief_dim
152        parquet_meta["final_belief_state"] = belief_data[-1].tolist()
153
154    table = pa.table(columns)
155    table = table.replace_schema_metadata({
156        **(table.schema.metadata or {}),
157        "am_metadata": json.dumps(parquet_meta),
158    })
159
160    pq.write_table(table, f"{file_prefix}.parquet", row_group_size=row_size )
161    save_json(parquet_meta, f"{file_prefix}.json")
162
163
164def minify_cpp(dfa: DFA, retain_names: bool = True) -> DFA:
165    state_list = list(dfa.states)
166    state_idx = {s: i for i, s in enumerate(state_list)}
167    symbol_list = sorted(dfa.input_symbols)
168    symbol_idx = {sym: i for i, sym in enumerate(symbol_list)}
169
170    adj = [[] for _ in range(len(state_list))]
171    for state, paths in dfa.transitions.items():
172        u = state_idx[state]
173        for sym, nxt in paths.items():
174            v = state_idx.get(nxt, -1)
175            if v != -1:
176                adj[u].append((symbol_idx[sym], v))
177
178    is_final = [s in dfa.final_states for s in state_list]
179    init_idx = state_idx[dfa.initial_state]
180
181    # result names in res now match the C++ struct fields exactly
182    res = minify_dfa_cpp(len(state_list), adj, init_idx, is_final)
183
184    if res.is_empty_language:
185        return dfa.__class__.empty_language(dfa.input_symbols)
186
187    if retain_names:
188        class_members = {}
189        for old_idx, nc in enumerate(res.eq_class):
190            if nc >= 0:
191                class_members.setdefault(nc, []).append(state_list[old_idx])
192        class_map = {nc: frozenset(m) for nc, m in class_members.items()}
193    else:
194        class_map = {nc: nc for nc in range(res.n_classes)}
195
196    new_states = set(class_map.values())
197    new_initial = class_map[res.new_initial]
198    new_final = {class_map[c] for c, f in enumerate(res.class_is_final) if f}
199
200    new_trans = {}
201    for nc in range(res.n_classes):
202        new_trans[class_map[nc]] = {
203            symbol_list[trans[0]]: class_map[trans[1]]
204            for trans in res.class_trans[nc]
205        }
206
207    return dfa.__class__(
208        states=new_states,
209        input_symbols=dfa.input_symbols,
210        transitions=new_trans,
211        initial_state=new_initial,
212        final_states=new_final,
213        allow_partial=any(len(t) < len(symbol_list) for t in new_trans.values())
214    )