mixedbench 1.0.0__py3-none-any.whl

mixedbench/__init__.py

@@ -0,0 +1 @@
+"""mixedbench: posterior approximation and prediction on real grouped data for GLMMs."""

mixedbench/contract/__init__.py

@@ -0,0 +1,21 @@
+from mixedbench.contract.materialised import contractHash, load, materialise
+from mixedbench.contract.problem import (
+    Prior,
+    Problem,
+    Refused,
+    Scale,
+    columnType,
+    withComposites,
+)
+
+__all__ = [
+    'Prior',
+    'Problem',
+    'Refused',
+    'Scale',
+    'columnType',
+    'contractHash',
+    'load',
+    'materialise',
+    'withComposites',
+]

mixedbench/contract/materialised.py

@@ -0,0 +1,124 @@
+"""materialised.py defines a problem's files on disk and the contract hash over them.
+
+A materialised problem is two files under the loader's cache:
+`table.parquet` holding the table, `problem.json` the record a table cannot and the contract hash.
+Three decisions guide them:
+- the files are the interface: `load` needs neither the loader nor the network,
+  so a suite user fetches the sources once per machine;
+- the hash is over values, not file bytes: a parquet writer's options or version do not move it;
+- the hash covers every field and every cell, the contract version among them,
+  so `load` refuses files that were edited or are stale.
+
+pipeline:
+---------
+a built Problem
+-> its record, every field but the table as JSON values     _record
+-> its contract hash                                        contractHash
+-> the table as parquet, the record and the hash as JSON    materialise
+-> the problem read back, refused if its hash moved         load
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import asdict, fields
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from mixedbench.contract.problem import Prior, Problem, Scale
+
+
+# =====================================================================================
+# the files
+# =====================================================================================
+# the package's interface to a problem: written once, loaded anywhere
+def materialise(problem: Problem, cache: Path) -> str:
+    """Write the problem's two files under `cache` and return its contract hash."""
+    dir = _problemDir(problem.id, cache)
+    dir.mkdir(parents=True, exist_ok=True)
+    problem.table.to_parquet(dir / 'table.parquet', index=False)
+    written = contractHash(problem)
+    stored = _record(problem) | {'hash': written}
+    (dir / 'problem.json').write_text(json.dumps(stored, sort_keys=True, indent=2))
+    return written
+
+
+def load(id: str, cache: Path) -> Problem:
+    """Raises where the files no longer give the hash they were written with (edited, or stale)."""
+    dir = _problemDir(id, cache)
+    stored = json.loads((dir / 'problem.json').read_text())
+    problem = Problem(
+        id=stored['id'],
+        table=pd.read_parquet(dir / 'table.parquet'),
+        design=tuple(stored['design']),
+        response=stored['response'],
+        factors=tuple(stored['factors']),
+        levels={column: tuple(levels) for column, levels in stored['levels'].items()},
+        scales={
+            regime: {column: Scale(**scale) for column, scale in columns.items()}
+            for regime, columns in stored['scales'].items()
+        },
+        dropped_share=stored['dropped_share'],
+        trials=stored['trials'],
+        family=stored['family'],
+        link=stored['link'],
+        prior={
+            name: Prior(prior['distribution'], tuple(prior['parameters']))
+            for name, prior in stored['prior'].items()
+        },
+        parameters=tuple(stored['parameters']),
+        correlations={factor: tuple(names) for factor, names in stored['correlations'].items()},
+        contract=stored['contract'],
+    )
+    if contractHash(problem) != stored['hash']:
+        raise ValueError(f'{id}: the files in {dir} do not match their contract hash; rebuild them')
+    return problem
+
+
+def _problemDir(id: str, cache: Path) -> Path:
+    """Named by the id's sha256: an object name may hold `/`, `!`, commas and spaces,
+    and the JSON carries the id itself."""
+    return Path(cache) / 'problems' / hashlib.sha256(id.encode()).hexdigest()[:16]
+
+
+# =====================================================================================
+# the contract hash
+# =====================================================================================
+# over values, so that a problem is the same whatever wrote its files
+def contractHash(problem: Problem) -> str:
+    """sha256 over the record with sorted keys,
+    then every column in table order as its name, dtype and little-endian bytes."""
+    canonical = json.dumps(_record(problem), sort_keys=True, separators=(',', ':'))
+    digest = hashlib.sha256(canonical.encode())
+    for name, column in problem.table.items():
+        dtype, values = _columnBytes(column)
+        digest.update(f'{name}\0{dtype}\0'.encode())
+        digest.update(values)
+    return digest.hexdigest()
+
+
+def _record(problem: Problem) -> dict:
+    """Every field but the table, as JSON values: a numpy level as its Python number."""
+    values = {f.name: getattr(problem, f.name) for f in fields(Problem) if f.name != 'table'}
+    return json.loads(json.dumps(values, default=_jsonValue))
+
+
+def _jsonValue(value: Scale | Prior | np.generic) -> dict | int | float:
+    if isinstance(value, np.generic):
+        return value.item()
+    return asdict(value)
+
+
+def _columnBytes(column: pd.Series) -> tuple[str, bytes]:
+    """A regime column as its codes, its categories in the dtype; a float as float64; an
+    integer in its own width."""
+    if isinstance(column.dtype, pd.CategoricalDtype):
+        categories = ','.join(column.cat.categories)
+        codes = column.cat.codes.to_numpy()
+        return f'category[{categories}]', codes.astype(codes.dtype.newbyteorder('<')).tobytes()
+    if pd.api.types.is_float_dtype(column.dtype):
+        return 'float64', column.to_numpy(dtype='<f8').tobytes()
+    values = column.to_numpy()
+    return str(values.dtype), values.astype(values.dtype.newbyteorder('<')).tobytes()

mixedbench/contract/problem.py

@@ -0,0 +1,412 @@
+"""problem.py defines the contract: a registry row and its table in, a `Problem` out.
+
+A problem is the table a model is fitted to and what it takes to fit and score it:
+the coded design and response, the family, the prior, the splits, the scales, the parameter names.
+`Problem.build` makes one from the typed row and the loader's typed table, by five decisions:
+- reproducible: a problem is a pure function of the pins, the row and the contract version;
+  its one random step, the splits, is seeded by the problem id;
+- uniform: the response type fixes the family, one prior serves every problem,
+  and one typing rule decides each column's kind (the row's documented prior is not used);
+- complete cases: a row missing a model value is dropped,
+  a row outside the contract (categorical response, censored model) raises `Refused`;
+- coded on the full table, scaled per regime: treatment contrasts against the first loaded level,
+  the table in raw units, each regime's moments recorded over the rows a fit sees;
+- brms's names in brms's block order, so draws and scores join on them.
+
+pipeline:
+---------
+a typed registry row and its loaded table
+-> refused if outside the contract             _refuseUncovered
+-> composites added, incomplete rows dropped   withComposites, _keptRows
+-> the coded design, response and factors      _kinds, _coded, _design, _codeFactor, _response
+-> the regimes `rows` and `groups`             _splits
+-> each regime's scales                        _scales
+-> the family, prior and parameter names       _prior, _parameters
+-> the problem                                 problem
+"""
+
+from __future__ import annotations
+
+import hashlib
+from dataclasses import dataclass
+from importlib.metadata import version
+
+import numpy as np
+import pandas as pd
+from mixedbench.registry.row import Row, parts
+
+# =====================================================================================
+# the problem
+# =====================================================================================
+# what `build` returns or raises, and the records a problem holds
+CONTRACT = int(version('mixedbench').partition('.')[0])  # the package's major version
+
+
+@dataclass(frozen=True)
+class Problem:
+    id: str  # <source>/<object>/<model_index>
+    table: pd.DataFrame  # the coded columns in raw units; one row per kept row
+    # column roles in `table`: the design columns in order, `Intercept` first; the response;
+    # the grouping factors as 0-based codes
+    design: tuple[str, ...]
+    response: str
+    factors: tuple[str, ...]
+    levels: dict[str, tuple]  # per recoded column, the level of each code (binary: 0 then 1)
+    # per regime (`full`, `rows`, `groups` where the problem has it), the moments of each
+    # continuous covariate and a continuous response over the rows a fit sees
+    scales: dict[str, dict[str, Scale]]
+    dropped_share: float  # of the loaded rows, those missing a model value
+    trials: str | None  # the trials column of a binomial response
+    family: str  # brms's names, by the family rule
+    link: str
+    prior: dict[str, Prior]  # per brms class the problem has
+    parameters: tuple[str, ...]  # the global parameters by the naming rule, in order
+    correlations: dict[str, tuple[str, ...]]  # per `cor[<factor>]` block, its elements in order
+    contract: int  # the version of the contract the problem was built under
+
+    @classmethod
+    def build(cls, row: Row, table: pd.DataFrame) -> Problem:
+        """Raises `Refused` where the contract does not cover the row."""
+        return _Builder(row, table).problem()
+
+
+@dataclass(frozen=True)
+class Prior:
+    distribution: str  # brms's name: normal, student_t (half on a scale), lkj
+    parameters: tuple[float, ...]
+
+
+@dataclass(frozen=True)
+class Scale:
+    mean: float
+    sd: float  # the sample standard deviation, as R's `scale`
+
+
+class Refused(Exception):
+    """The contract states no problem for the row: the message names the row, `reason` says
+    why."""
+
+    def __init__(self, id: str, reason: str) -> None:
+        super().__init__(f'{id}: {reason}')
+        self.reason = reason
+
+
+# =====================================================================================
+# the builder
+# =====================================================================================
+# the pipeline behind `Problem.build`, one step per rule of the contract
+Coded = dict[str, dict[str, pd.Series]]  # per covariate or interaction, its coded columns
+
+
+class _Builder:
+    # the family rule: by response type, brms's family and link
+    FAMILY = {
+        'continuous': ('gaussian', 'identity'),
+        'binary': ('bernoulli', 'logit'),
+        'count': ('poisson', 'log'),
+        'binomial': ('binomial', 'logit'),
+        'ordinal': ('cumulative', 'logit'),
+    }
+    CENSORED = ('cox', 'weibull')
+    # the prior: on the standardised scale, per brms class
+    PRIOR = {
+        'Intercept': Prior('normal', (0, 2.5)),
+        'b': Prior('normal', (0, 2.5)),
+        'sd': Prior('student_t', (3, 0, 2.5)),
+        'cor': Prior('lkj', (1,)),
+        'sigma': Prior('student_t', (3, 0, 2.5)),
+    }
+    # the splits: two regimes, each a column of `seen` and `held`
+    REGIME = pd.CategoricalDtype(['seen', 'held'])
+
+    def __init__(self, row: Row, table: pd.DataFrame) -> None:
+        self.row = row
+        self.id = row.id
+        self._refuseUncovered()
+        self.table = self._keptRows(withComposites(table, row.grouping_factors))
+        if self.table.empty:  # nlme Phenobarb: doses and responses on different rows
+            raise Refused(self.id, 'no row has every model value')
+        self.dropped_share = (len(table) - len(self.table)) / len(table)
+
+    def problem(self) -> Problem:
+        kinds = self._kinds()
+        coded = self._coded(kinds)
+        design = self._design(coded)
+        factors = {factor: self._codeFactor(factor) for factor in self.row.grouping_factors}
+        regimes = self._splits(factors)
+        table = pd.DataFrame(design | self._response() | factors | regimes).reset_index(drop=True)
+        correlations = self._correlations(coded)
+        family, link = self.FAMILY[self.row.response_type]
+        return Problem(
+            id=self.id,
+            table=table,
+            design=tuple(design),
+            response=self.row.response,
+            factors=self.row.grouping_factors,
+            levels=self._levels(kinds),
+            scales=self._scales(table, kinds),
+            dropped_share=self.dropped_share,
+            trials=self.row.trials or None,
+            family=family,
+            link=link,
+            prior=self._prior(design, correlations),
+            parameters=self._parameters(design, coded, correlations),
+            correlations=correlations,
+            contract=CONTRACT,
+        )
+
+    # -------------------------------------------------------------------------------------
+    # rows -> coding -> splits -> scales -> model
+    # -------------------------------------------------------------------------------------
+
+    def _refuseUncovered(self) -> None:
+        """A categorical response and a censored model are outside suite 1.1; every other
+        `family` takes the family rule."""
+        if self.row.response_type == 'categorical':
+            raise Refused(self.id, 'a categorical response is outside suite 1.1 (ADR 0001)')
+        if self.row.family in self.CENSORED:
+            reason = f'family {self.row.family} is a censored model, outside suite 1.1 (ADR 0001)'
+            raise Refused(self.id, reason)
+
+    def _keptRows(self, table: pd.DataFrame) -> pd.DataFrame:
+        """A row with a missing value in any model column is dropped; a composite factor is
+        missing where a part is."""
+        model = [self.row.response, *self.row.covariates]
+        for factor in self.row.grouping_factors:
+            model += parts(factor)
+        if self.row.trials:
+            model.append(self.row.trials)
+        return table.dropna(subset=model)
+
+    def _kinds(self) -> dict[str, str]:
+        """Per covariate, its kind by the typing rule."""
+        kinds = {}
+        for column in self.row.covariates:
+            kinds[column], _ = columnType(self.table[column], self.row.types.get(column, ''))
+        return kinds
+
+    def _coded(self, kinds: dict[str, str]) -> Coded:
+        """Per covariate, then per interaction, its coded columns."""
+        coded = {}
+        for covariate, kind in kinds.items():
+            coded[covariate] = self._codeCovariate(covariate, kind)
+        for term in self.row.interactions:
+            coded[term] = self._codeInteraction(term, coded)
+        return coded
+
+    def _codeCovariate(self, column: str, kind: str) -> dict[str, pd.Series]:
+        """Continuous raw; binary 0/1 with the first loaded level 0; a categorical or ordered
+        factor as K − 1 treatment contrasts against the first level, `<column>[<level>]`."""
+        values = self.table[column]
+        if kind == 'continuous':
+            return {column: values.astype(float)}
+        loaded = self._levelOrder(values)
+        if kind == 'binary':
+            return {column: (values != loaded[0]).astype(float)}
+        return {f'{column}[{level}]': (values == level).astype(float) for level in loaded[1:]}
+
+    @staticmethod
+    def _codeInteraction(term: str, coded: Coded) -> dict[str, pd.Series]:
+        """The products of the coded columns of its parts, `<a>:<b>` per pair in order."""
+        columns = {'': 1.0}
+        for part in parts(term):
+            products = {}
+            for name, column in columns.items():
+                for partName, partColumn in coded[part].items():
+                    products[f'{name}:{partName}'.removeprefix(':')] = column * partColumn
+            columns = products
+        return columns
+
+    def _design(self, coded: Coded) -> dict[str, pd.Series]:
+        """`Intercept` first, then the coded columns in the row's order."""
+        design = {'Intercept': pd.Series(1.0, index=self.table.index)}
+        for columns in coded.values():
+            design.update(columns)
+        return design
+
+    def _codeFactor(self, factor: str) -> pd.Series:
+        values = self.table[factor]
+        codes = pd.Categorical(values, categories=self._levelOrder(values)).codes
+        return pd.Series(codes, index=self.table.index)
+
+    def _response(self) -> dict[str, pd.Series]:
+        """By `response_type`: continuous, count and binomial successes as loaded, binary 0/1
+        with the first loaded level 0, ordinal its 1-based index in the loaded level order; a
+        binomial's trials beside it."""
+        values = self.table[self.row.response]
+        if self.row.response_type == 'binary':
+            values = (values != self._levelOrder(values)[0]).astype(int)
+        if self.row.response_type == 'ordinal':
+            codes = pd.Categorical(values, categories=self._levelOrder(values)).codes
+            values = pd.Series(codes + 1, index=self.table.index)
+        response = {self.row.response: values}
+        if self.row.trials:
+            response[self.row.trials] = self.table[self.row.trials]
+        return response
+
+    def _levels(self, kinds: dict[str, str]) -> dict[str, tuple]:
+        """The loaded level order of every recoded column: a covariate that is not continuous,
+        a binary or ordinal response, a grouping factor."""
+        recoded = [column for column, kind in kinds.items() if kind != 'continuous']
+        if self.row.response_type in ('binary', 'ordinal'):
+            recoded.append(self.row.response)
+        recoded += self.row.grouping_factors
+        return {column: self._levelOrder(self.table[column]) for column in recoded}
+
+    def _splits(self, factors: dict[str, pd.Series]) -> dict[str, pd.Series]:
+        """Seeded by the first 8 bytes of sha256 of the id; rows draw before groups. A first
+        factor of two groups has no `groups` regime: holding one would leave one."""
+        seed = int.from_bytes(hashlib.sha256(self.id.encode()).digest()[:8], 'big')
+        rng = np.random.default_rng(seed)
+        regimes = {'rows': self._regime(self._heldRows(rng, factors))}
+        first = factors[self.row.grouping_factors[0]]
+        if first.max() + 1 > 2:
+            regimes['groups'] = self._regime(self._heldGroups(rng, first))
+        return regimes
+
+    def _heldRows(self, rng: np.random.Generator, factors: dict[str, pd.Series]) -> np.ndarray:
+        """Every group of every factor keeps seen the row it shows first in the seeded order;
+        of the other rows, the first nearest whole number to 10% of all rows are held."""
+        n = len(self.table)
+        order = self._seededOrder(rng, n)  # (n,) row positions
+        spare = np.ones(n, dtype=bool)  # (n,) in seeded order
+        for codes in factors.values():
+            spare &= pd.Series(codes.to_numpy()[order]).duplicated().to_numpy()
+        held = np.zeros(n, dtype=bool)  # (n,)
+        held[order[spare][: (n + 5) // 10]] = True
+        return held
+
+    def _heldGroups(self, rng: np.random.Generator, codes: pd.Series) -> np.ndarray:
+        """The nearest whole number to 10% of the groups, at least one, at least two left."""
+        groups = codes.max() + 1
+        count = min(max(1, (groups + 5) // 10), groups - 2)
+        held = self._seededOrder(rng, groups)[:count]
+        return np.isin(codes.to_numpy(), held)  # (n,)
+
+    @staticmethod
+    def _seededOrder(rng: np.random.Generator, n: int) -> np.ndarray:
+        """An argsort of uniform draws, whose stream numpy keeps across versions (a shuffle's
+        algorithm it does not promise)."""
+        return np.argsort(rng.random(n), kind='stable')  # (n,)
+
+    def _regime(self, held: np.ndarray) -> pd.Series:
+        regime = pd.Categorical.from_codes(held.astype(int), dtype=self.REGIME)
+        return pd.Series(regime, index=self.table.index)
+
+    def _scales(self, table: pd.DataFrame, kinds: dict[str, str]) -> dict[str, dict[str, Scale]]:
+        """The standardisation per regime on the rows a fit sees; the table stays in raw units."""
+        columns = [column for column, kind in kinds.items() if kind == 'continuous']
+        if self.row.response_type == 'continuous':
+            columns.append(self.row.response)
+        seen = {'full': pd.Series(True, index=table.index)}
+        for regime in ('rows', 'groups'):
+            if regime in table:
+                seen[regime] = table[regime] == 'seen'
+        scales = {}
+        for regime, rows in seen.items():
+            values = table.loc[rows, columns]
+            scales[regime] = {
+                c: Scale(float(values[c].mean()), float(values[c].std())) for c in columns
+            }
+        return scales
+
+    def _correlations(self, coded: Coded) -> dict[str, tuple[str, ...]]:
+        """Per factor with a correlated slope, the elements of `cor[<factor>]`: `Intercept`,
+        then the coded columns of its correlated slopes in order."""
+        correlations = {}
+        for term in self.row.random_slopes:
+            if term.correlated:
+                elements = correlations.get(term.factor, ('Intercept',))
+                correlations[term.factor] = elements + tuple(coded[term.covariate])
+        return correlations
+
+    def _prior(
+        self, design: dict[str, pd.Series], correlations: dict[str, tuple[str, ...]]
+    ) -> dict[str, Prior]:
+        """The prior over the classes the problem has (brms refuses one on an absent
+        class): no `Intercept` on an ordinal model, whose thresholds take brms's default."""
+        classes = []
+        if self.row.response_type != 'ordinal':
+            classes.append('Intercept')
+        if len(design) > 1:
+            classes.append('b')
+        classes.append('sd')
+        if correlations:
+            classes.append('cor')
+        if self.row.response_type == 'continuous':
+            classes.append('sigma')
+        return {name: self.PRIOR[name] for name in classes}
+
+    def _parameters(
+        self,
+        design: dict[str, pd.Series],
+        coded: Coded,
+        correlations: dict[str, tuple[str, ...]],
+    ) -> tuple[str, ...]:
+        """brms's block order: the fixed effects (an ordinal model's thresholds `Intercept[k]`
+        in place of `Intercept`, brms's name without its `b_`), per factor `sd[<factor>]`, the
+        slopes `sd[<x>|<factor>]` and `cor[<factor>]` where a slope is correlated, `sigma`."""
+        names = list(design)
+        if self.row.response_type == 'ordinal':
+            thresholds = len(self._levelOrder(self.table[self.row.response])) - 1
+            names = [f'Intercept[{k}]' for k in range(1, thresholds + 1)] + names[1:]
+        for factor in self.row.grouping_factors:
+            names += self._randomEffects(factor, coded)
+            if factor in correlations:
+                names.append(f'cor[{factor}]')
+        if self.row.response_type == 'continuous':
+            names.append('sigma')
+        return tuple(names)
+
+    def _randomEffects(self, factor: str, coded: Coded) -> list[str]:
+        names = [f'sd[{factor}]']
+        for term in self.row.random_slopes:
+            if term.factor == factor:
+                names += [f'sd[{slope}|{factor}]' for slope in coded[term.covariate]]
+        return names
+
+    @staticmethod
+    def _levelOrder(values: pd.Series) -> tuple:
+        """The loaded level order: a categorical's categories present in the column, else the
+        sorted distinct values (what the loader gives a plain column declared categorical)."""
+        if isinstance(values.dtype, pd.CategoricalDtype):
+            return tuple(values.cat.remove_unused_categories().cat.categories)
+        return tuple(sorted(values.dropna().unique()))
+
+
+# =====================================================================================
+# the rules shared with the survey
+# =====================================================================================
+# the survey measures a row's design width with them, so it is the width a build gives
+MAX_INTEGER_LEVELS = 10
+
+
+def columnType(values: pd.Series, declared: str) -> tuple[str, int | str]:
+    """A column's kind and its level count K (blank if continuous); the first matching rule wins."""
+    distinct = values.dropna().nunique()
+    dtype = values.dtype
+    if isinstance(dtype, pd.CategoricalDtype) and dtype.ordered:
+        return 'ordered', distinct
+    if distinct == 2:
+        return 'binary', 2
+    if declared == 'continuous':
+        return 'continuous', ''
+    if isinstance(dtype, pd.CategoricalDtype) or pd.api.types.is_string_dtype(dtype):
+        return 'categorical', distinct
+    if pd.api.types.is_integer_dtype(dtype) and distinct <= MAX_INTEGER_LEVELS:
+        return 'categorical', distinct
+    return 'continuous', ''
+
+
+def withComposites(table: pd.DataFrame, factors: tuple[str, ...]) -> pd.DataFrame:
+    """Adds each composite factor `a:b` as a column of its parts' joined levels;
+    a row missing a part has none."""
+    composites = {}
+    for factor in factors:
+        columns = parts(factor)
+        if len(columns) > 1:
+            composite = table[columns[0]].astype(str)
+            for part in columns[1:]:
+                composite = composite + ':' + table[part].astype(str)  # missing stays missing
+            composites[factor] = composite
+    return table.assign(**composites)

mixedbench/datasets/__init__.py

@@ -0,0 +1,11 @@
+from mixedbench.datasets.loader import (
+    ColumnType,
+    DataObject,
+    Loader,
+    Manifest,
+    ManifestRow,
+    Source,
+    fetch,
+)
+
+__all__ = ['ColumnType', 'DataObject', 'Loader', 'Manifest', 'ManifestRow', 'Source', 'fetch']