|
| 1 | +import numpy as np |
| 2 | + |
| 3 | +from pymc import Model |
| 4 | +from pymc.printing import str_for_dist, str_for_potential_or_deterministic |
| 5 | +from pytensor import Mode |
| 6 | +from pytensor.compile.sharedvalue import SharedVariable |
| 7 | +from pytensor.graph.type import Constant, Variable |
| 8 | +from rich.box import SIMPLE_HEAD |
| 9 | +from rich.table import Table |
| 10 | + |
| 11 | + |
| 12 | +def variable_expression( |
| 13 | + model: Model, |
| 14 | + var: Variable, |
| 15 | + truncate_deterministic: int | None, |
| 16 | +) -> str: |
| 17 | + """Get the expression of a variable in a human-readable format.""" |
| 18 | + if var in model.data_vars: |
| 19 | + var_expr = "Data" |
| 20 | + elif var in model.deterministics: |
| 21 | + str_repr = str_for_potential_or_deterministic(var, dist_name="") |
| 22 | + _, var_expr = str_repr.split(" ~ ") |
| 23 | + var_expr = var_expr[1:-1] # Remove outer parentheses (f(...)) |
| 24 | + if truncate_deterministic is not None and len(var_expr) > truncate_deterministic: |
| 25 | + contents = var_expr[2:-1].split(", ") |
| 26 | + str_len = 0 |
| 27 | + for show_n, content in enumerate(contents): |
| 28 | + str_len += len(content) + 2 |
| 29 | + if str_len > truncate_deterministic: |
| 30 | + break |
| 31 | + var_expr = f"f({', '.join(contents[:show_n])}, ...)" |
| 32 | + elif var in model.potentials: |
| 33 | + var_expr = str_for_potential_or_deterministic(var, dist_name="Potential").split(" ~ ")[1] |
| 34 | + else: # basic_RVs |
| 35 | + var_expr = str_for_dist(var).split(" ~ ")[1] |
| 36 | + return var_expr |
| 37 | + |
| 38 | + |
| 39 | +def _extract_dim_value(var: SharedVariable | Constant) -> np.ndarray: |
| 40 | + if isinstance(var, SharedVariable): |
| 41 | + return var.get_value(borrow=True) |
| 42 | + else: |
| 43 | + return var.data |
| 44 | + |
| 45 | + |
| 46 | +def dims_expression(model: Model, var: Variable) -> str: |
| 47 | + """Get the dimensions of a variable in a human-readable format.""" |
| 48 | + if (dims := model.named_vars_to_dims.get(var.name)) is not None: |
| 49 | + dim_sizes = {dim: _extract_dim_value(model.dim_lengths[dim]) for dim in dims} |
| 50 | + return " × ".join(f"{dim}[{dim_size}]" for dim, dim_size in dim_sizes.items()) |
| 51 | + else: |
| 52 | + dim_sizes = list(var.shape.eval(mode=Mode(linker="py", optimizer="fast_compile"))) |
| 53 | + return f"[{', '.join(map(str, dim_sizes))}]" if dim_sizes else "" |
| 54 | + |
| 55 | + |
| 56 | +def model_parameter_count(model: Model) -> int: |
| 57 | + """Count the number of parameters in the model.""" |
| 58 | + rv_shapes = model.eval_rv_shapes() # Includes transformed variables |
| 59 | + return np.sum([np.prod(rv_shapes[free_rv.name]).astype(int) for free_rv in model.free_RVs]) |
| 60 | + |
| 61 | + |
| 62 | +def model_table( |
| 63 | + model: Model, |
| 64 | + *, |
| 65 | + split_groups: bool = True, |
| 66 | + truncate_deterministic: int | None = None, |
| 67 | + parameter_count: bool = True, |
| 68 | +) -> Table: |
| 69 | + """Create a rich table with a summary of the model's variables and their expressions. |
| 70 | +
|
| 71 | + Parameters |
| 72 | + ---------- |
| 73 | + model : Model |
| 74 | + The PyMC model to summarize. |
| 75 | + split_groups : bool |
| 76 | + If True, each group of variables (data, free_RVs, deterministics, potentials, observed_RVs) |
| 77 | + will be separated by a section. |
| 78 | + truncate_deterministic : int | None |
| 79 | + If not None, truncate the expression of deterministic variables that go beyond this length. |
| 80 | + empty_dims : bool |
| 81 | + If True, show the dimensions of scalar variables as an empty list. |
| 82 | + parameter_count : bool |
| 83 | + If True, add a row with the total number of parameters in the model. |
| 84 | +
|
| 85 | + Returns |
| 86 | + ------- |
| 87 | + Table |
| 88 | + A rich table with the model's variables, their expressions and dims. |
| 89 | +
|
| 90 | + Examples |
| 91 | + -------- |
| 92 | + .. code-block:: python |
| 93 | +
|
| 94 | + import numpy as np |
| 95 | + import pymc as pm |
| 96 | +
|
| 97 | + from pymc_experimental.printing import model_table |
| 98 | +
|
| 99 | + coords = {"subject": range(20), "param": ["a", "b"]} |
| 100 | + with pm.Model(coords=coords) as m: |
| 101 | + x = pm.Data("x", np.random.normal(size=(20, 2)), dims=("subject", "param")) |
| 102 | + y = pm.Data("y", np.random.normal(size=(20,)), dims="subject") |
| 103 | +
|
| 104 | + beta = pm.Normal("beta", mu=0, sigma=1, dims="param") |
| 105 | + mu = pm.Deterministic("mu", pm.math.dot(x, beta), dims="subject") |
| 106 | + sigma = pm.HalfNormal("sigma", sigma=1) |
| 107 | +
|
| 108 | + y_obs = pm.Normal("y_obs", mu=mu, sigma=sigma, observed=y, dims="subject") |
| 109 | +
|
| 110 | + table = model_table(m) |
| 111 | + table # Displays the following table in an interactive environment |
| 112 | + ''' |
| 113 | + Variable Expression Dimensions |
| 114 | + ───────────────────────────────────────────────────── |
| 115 | + x = Data subject[20] × param[2] |
| 116 | + y = Data subject[20] |
| 117 | +
|
| 118 | + beta ~ Normal(0, 1) param[2] |
| 119 | + sigma ~ HalfNormal(0, 1) |
| 120 | + Parameter count = 3 |
| 121 | +
|
| 122 | + mu = f(beta) subject[20] |
| 123 | +
|
| 124 | + y_obs ~ Normal(mu, sigma) subject[20] |
| 125 | + ''' |
| 126 | +
|
| 127 | + Output can be explicitly rendered in a rich console or exported to text, html or svg. |
| 128 | +
|
| 129 | + .. code-block:: python |
| 130 | +
|
| 131 | + from rich.console import Console |
| 132 | +
|
| 133 | + console = Console(record=True) |
| 134 | + console.print(table) |
| 135 | + text_export = console.export_text() |
| 136 | + html_export = console.export_html() |
| 137 | + svg_export = console.export_svg() |
| 138 | +
|
| 139 | + """ |
| 140 | + table = Table( |
| 141 | + show_header=True, |
| 142 | + show_edge=False, |
| 143 | + box=SIMPLE_HEAD, |
| 144 | + highlight=False, |
| 145 | + collapse_padding=True, |
| 146 | + ) |
| 147 | + table.add_column("Variable", justify="right") |
| 148 | + table.add_column("Expression", justify="left") |
| 149 | + table.add_column("Dimensions") |
| 150 | + |
| 151 | + if split_groups: |
| 152 | + groups = ( |
| 153 | + model.data_vars, |
| 154 | + model.free_RVs, |
| 155 | + model.deterministics, |
| 156 | + model.potentials, |
| 157 | + model.observed_RVs, |
| 158 | + ) |
| 159 | + else: |
| 160 | + # Show variables in the order they were defined |
| 161 | + groups = (model.named_vars.values(),) |
| 162 | + |
| 163 | + for group in groups: |
| 164 | + if not group: |
| 165 | + continue |
| 166 | + |
| 167 | + for var in group: |
| 168 | + var_name = var.name |
| 169 | + sep = f'[b]{" ~" if (var in model.basic_RVs) else " ="}[/b]' |
| 170 | + var_expr = variable_expression(model, var, truncate_deterministic) |
| 171 | + dims_expr = dims_expression(model, var) |
| 172 | + if dims_expr == "[]": |
| 173 | + dims_expr = "" |
| 174 | + table.add_row(var_name + sep, var_expr, dims_expr) |
| 175 | + |
| 176 | + if parameter_count and (not split_groups or group == model.free_RVs): |
| 177 | + n_parameters = model_parameter_count(model) |
| 178 | + table.add_row("", "", f"[i]Parameter count = {n_parameters}[/i]") |
| 179 | + |
| 180 | + table.add_section() |
| 181 | + |
| 182 | + return table |
0 commit comments