Updated example data

groupselect/examples/__init__.py

@@ -1,6 +1,12 @@
-"""Example participant datasets loaded from CSV files for demos and tests."""
+"""Example participant datasets, loaded from CSV or generated on the fly."""
 
+from collections.abc import Mapping
 from pathlib import Path
+from typing import Iterable
+
+import numpy as np
+
+from groupselect.field_mode import FieldMode
 
 try:
     import pandas as pd
@@ -14,6 +20,28 @@
 DATA_DIR = Path(__file__).parent / "data"
 
 
+class _MissingPandasExampleData(Mapping):
+    """Stand-in for the example-data dicts when pandas isn't installed.
+
+    The example datasets are stored as CSV files and only loaded via
+    pandas, so there is nothing to fall back to without it. Any lookup
+    raises a clear error rather than a plain ``KeyError``.
+    """
+
+    def __getitem__(self, key):
+        raise ImportError(
+            f"Example dataset {key!r} is not available: pandas is not "
+            "installed. Install pandas to use groupselect.examples "
+            "example_data_np / example_data_pd."
+        )
+
+    def __iter__(self):
+        return iter(())
+
+    def __len__(self):
+        return 0
+
+
 if HAS_PANDAS:
     example_data_np = {}
     example_data_pd = {}
@@ -31,5 +59,139 @@
             .to_numpy()
         )
 else:
-    # TODO: Load example data when pandas is not available.
-    example_data_np = {}
+    example_data_np = _MissingPandasExampleData()
+    example_data_pd = _MissingPandasExampleData()
+
+
+# ---------------------------------------------------------------------------
+# Synthetic dataset generation.
+# ---------------------------------------------------------------------------
+
+
+def generate_participants(
+    n_participants: int,
+    n_div_fields: int,
+    n_field_features: int | Iterable[int] = 3,
+    rng: np.random.Generator | int | None = None,
+) -> np.ndarray:
+    """Generate a synthetic array of participants for benchmarking.
+
+    Each of the ``n_div_fields`` columns is populated independently:
+    feature codes ``0 .. n_field_features - 1`` are assigned so that
+    every feature occurs an (as close to) equal number of times
+    within that field, then the column is randomly shuffled. Passing
+    a single int for ``n_field_features`` gives every field the same
+    number of features; passing a sequence of length ``n_div_fields``
+    allows each field to have a different number of features.
+
+    Returns a ``(n_participants, n_div_fields)`` array of ``int``
+    feature codes, directly usable as the ``participants`` argument
+    of :func:`groupselect.allocate_numpy`.
+    """
+    if n_participants <= 0:
+        raise ValueError("n_participants must be a positive integer.")
+    if n_div_fields <= 0:
+        raise ValueError("n_div_fields must be a positive integer.")
+
+    if isinstance(n_field_features, int):
+        n_field_features = [n_field_features] * n_div_fields
+    else:
+        n_field_features = list(n_field_features)
+    if len(n_field_features) != n_div_fields:
+        raise ValueError(
+            "n_field_features must be an int, or a sequence of "
+            "length n_div_fields."
+        )
+    if any(n_features < 1 for n_features in n_field_features):
+        raise ValueError("n_field_features must be at least 1.")
+
+    rng = (
+        rng
+        if isinstance(rng, np.random.Generator)
+        else np.random.default_rng(rng)
+    )
+
+    columns = []
+    for n_features in n_field_features:
+        # Distribute feature codes as evenly as possible over the
+        # participants, then shuffle so codes aren't grouped by index.
+        column = np.arange(n_participants) % n_features
+        rng.shuffle(column)
+        columns.append(column)
+
+    return np.column_stack(columns).astype(int)
+
+
+def generate_fields(
+    n_div_fields: int, mode: FieldMode = FieldMode.Diversify
+) -> dict[int, FieldMode]:
+    """Return a fields dict marking columns ``0 .. n_div_fields - 1``."""
+    return {field_id: mode for field_id in range(n_div_fields)}
+
+
+def generate_dataset(
+    n_participants: int,
+    n_div_fields: int,
+    n_field_features: int | Iterable[int] = 3,
+    n_part_per_group: int = 6,
+    n_allocations: int = 1,
+    rng: np.random.Generator | int | None = None,
+) -> dict:
+    """Generate a full synthetic dataset ready for `allocate_numpy`.
+
+    Returns a dict with keys ``participants``, ``fields`` and
+    ``n_part_per_group``, which can be passed straight into
+    :func:`groupselect.allocate_numpy` via ``**``, e.g.::
+
+        dataset = generate_dataset(
+            n_participants=60, n_div_fields=3, n_field_features=4,
+            n_part_per_group=10, n_allocations=3,
+        )
+        result = groupselect.allocate_numpy(
+            **dataset, algorithm="HERMES", settings=...,
+        )
+    """
+    return dict(
+        participants=generate_participants(
+            n_participants=n_participants,
+            n_div_fields=n_div_fields,
+            n_field_features=n_field_features,
+            rng=rng,
+        ),
+        fields=generate_fields(n_div_fields),
+        n_part_per_group=n_allocations * [n_part_per_group],
+    )
+
+
+if HAS_PANDAS:
+
+    def generate_dataset_pandas(
+        n_participants: int,
+        n_div_fields: int,
+        n_field_features: int | Iterable[int] = 3,
+        n_part_per_group: int = 6,
+        n_allocations: int = 1,
+        rng: np.random.Generator | int | None = None,
+    ) -> dict:
+        """Like `generate_dataset`, but with participants as a DataFrame.
+
+        The DataFrame has diversity-field columns named ``"field_0"``,
+        ``"field_1"``, ... and an ``"ID"``-named index, ready for
+        :func:`groupselect.allocate_pandas`.
+        """
+        dataset = generate_dataset(
+            n_participants=n_participants,
+            n_div_fields=n_div_fields,
+            n_field_features=n_field_features,
+            n_part_per_group=n_part_per_group,
+            n_allocations=n_allocations,
+            rng=rng,
+        )
+        field_names = [f"field_{i}" for i in range(n_div_fields)]
+        participants = pd.DataFrame(
+            dataset["participants"], columns=field_names
+        )
+        participants.index.name = "ID"
+        dataset["participants"] = participants
+        dataset["fields"] = dict(zip(field_names, dataset["fields"].values()))
+        return dataset

groupselect/examples/data/30i_4j_3k_3d_data.csv

@@ -1,31 +1,31 @@
-,ID,gender,age,region,fixed,allocation
-0,1,male,older,rural,False,0
-1,2,male,younger,urban,False,0
-2,3,female,younger,rural,False,0
-3,4,female,older,urban,False,0
-4,5,male,younger,urban,False,0
-5,6,male,older,rural,False,0
-6,7,male,older,rural,False,0
-7,8,female,older,rural,False,1
-8,9,male,younger,urban,False,1
-9,10,female,younger,urban,False,1
-10,11,male,older,rural,False,1
-11,12,female,older,rural,False,1
-12,13,male,older,urban,False,1
-13,14,female,younger,urban,False,1
-14,15,female,older,urban,False,2
-15,16,female,younger,urban,False,2
-16,17,male,younger,urban,False,2
-17,18,female,older,urban,False,2
-18,19,female,older,urban,False,2
-19,20,male,older,urban,False,2
-20,21,male,younger,urban,False,2
-21,22,male,younger,urban,False,2
-22,23,female,older,urban,False,3
-23,24,male,older,urban,False,3
-24,25,female,younger,rural,False,3
-25,26,male,younger,rural,False,3
-26,27,female,older,urban,False,3
-27,28,male,older,urban,False,3
-28,29,female,younger,urban,False,3
-29,30,male,younger,urban,False,3
+ID,gender,age,region,fixed,allocation
+1,male,older,rural,False,0
+2,male,younger,urban,False,0
+3,female,younger,rural,False,0
+4,female,older,urban,False,0
+5,male,younger,urban,False,0
+6,male,older,rural,False,0
+7,male,older,rural,False,0
+8,female,older,rural,False,1
+9,male,younger,urban,False,1
+10,female,younger,urban,False,1
+11,male,older,rural,False,1
+12,female,older,rural,False,1
+13,male,older,urban,False,1
+14,female,younger,urban,False,1
+15,female,older,urban,False,2
+16,female,younger,urban,False,2
+17,male,younger,urban,False,2
+18,female,older,urban,False,2
+19,female,older,urban,False,2
+20,male,older,urban,False,2
+21,male,younger,urban,False,2
+22,male,younger,urban,False,2
+23,female,older,urban,False,3
+24,male,older,urban,False,3
+25,female,younger,rural,False,3
+26,male,younger,rural,False,3
+27,female,older,urban,False,3
+28,male,older,urban,False,3
+29,female,younger,urban,False,3
+30,male,younger,urban,False,3

groupselect/examples/data/default.csv

@@ -0,0 +1,111 @@
+name,dietary reqs,gender,age,education,ethnicity,region,photo consent,audio consent
+Abner Nigel,no fish,Female,45-59,Level 4 and above,White,Urban,yes,yes
+Lane Hayward,normal,Male,16-29,No Qualifications/Level 1,White,Urban,yes,yes
+Jayde Patterson,vegetarian,Female,45-59,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Eleanora Peterson,vegetarian,Female,60+,No Qualifications/Level 1,White,Urban,yes,yes
+Wilson Selby,vegetarian,Male,16-29,No Qualifications/Level 1,White,Urban,yes,yes
+Caelie Abrams,normal,Male,16-29,No Qualifications/Level 1,White,Rural,maybe,yes
+Brooklyn Waller,normal,Female,45-59,Level 2/Level 3/Apprenticeship/Other,White,Urban,maybe,yes
+Jemima Hall,vegetarian,Female,30-44,Level 2/Level 3/Apprenticeship/Other,White,Urban,maybe,no
+Rae Hooper,normal,Male,60+,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,no
+Alannis Porter,normal,Male,16-29,No Qualifications/Level 1,White,Urban,yes,yes
+Jerald Cropper,vegetarian,Female,30-44,Level 4 and above,BAME,Urban,yes,yes
+Richardine Mercer,normal,Female,45-59,Level 2/Level 3/Apprenticeship/Other,White,Urban,maybe,yes
+Louie Roydon,normal,Male,60+,No Qualifications/Level 1,BAME,Urban,yes,yes
+Albin Albertson,normal,Male,60+,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Jaynie Hillam,normal,Female,45-59,No Qualifications/Level 1,BAME,Urban,yes,yes
+Darnell Marley,normal,Female,30-44,Level 2/Level 3/Apprenticeship/Other,BAME,Urban,yes,yes
+Adriana Wardrobe,vegetarian,Male,60+,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Rebecca Taylor,normal,Female,45-59,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,no
+Lavonne Payton,normal,Male,60+,Level 2/Level 3/Apprenticeship/Other,BAME,Urban,yes,yes
+Keri Sutton,normal,Male,60+,Level 4 and above,White,Urban,no,yes
+Jemmy Forrest,normal,Female,60+,No Qualifications/Level 1,White,Urban,yes,yes
+Kaety Woods,normal,Male,45-59,No Qualifications/Level 1,White,Urban,yes,yes
+Marni Shepherd,normal,Male,45-59,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,no
+Meriel Eccleston,vegetarian,Male,16-29,No Qualifications/Level 1,w,Rural,yes,yes
+Morty Thwaite,normal,Male,16-29,Level 2/Level 3/Apprenticeship/Other,BAME,Urban,yes,yes
+Zack Woodham,normal,Male,30-44,Level 4 and above,White,Urban,yes,yes
+Eddie Christopher,normal,Male,45-59,No Qualifications/Level 1,White,Urban,yes,yes
+Beaumont Dickson,normal,Female,30-44,Level 4 and above,White,Rural,yes,yes
+Zac Geary,normal,Male,45-59,No Qualifications/Level 1,White,Rural,yes,yes
+Suzan Wilmer,normal,Female,16-29,No Qualifications/Level 1,White,Urban,yes,yes
+Mia Battle,normal,Female,60+,Level 4 and above,White,Urban,maybe,yes
+Krystine Brassington,normal,Male,16-29,No Qualifications/Level 1,White,Urban,yes,yes
+Malcolm Bell,nut allergy,Male,60+,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Tamia Lum,normal,Female,16-29,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Korey Bone,normal,Male,30-44,Level 4 and above,White,Rural,no,no
+Patton Jacobs,vegetarian,Male,45-59,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Garrett Powers,normal,Female,60+,No Qualifications/Level 1,White,Urban,yes,yes
+Delight Blackbourne,normal,Female,60+,No Qualifications/Level 1,White,Urban,maybe,yes
+Cal Hodges,normal,Male,30-44,No Qualifications/Level 1,White,Urban,yes,yes
+Nellie Denzil,normal,Male,30-44,Level 4 and above,White,Urban,maybe,yes
+Don Bannister,normal,Male,30-44,Level 2/Level 3/Apprenticeship/Other,White,Urban,maybe,yes
+Sheard Lovell,normal,Male,60+,Level 4 and above,BAME,Urban,yes,yes
+Theodora Strudwick,vegetarian,Male,16-29,Level 4 and above,BAME,Urban,yes,yes
+Lennie Myers,normal,Male,60+,Level 4 and above,White,Urban,yes,yes
+Dene Cox,normal,Male,30-44,No Qualifications/Level 1,BAME,Urban,maybe,yes
+Marley Waterman,normal,Male,60+,Level 2/Level 3/Apprenticeship/Other,White,Rural,maybe,yes
+Meadow Southgate,normal,Female,45-59,No Qualifications/Level 1,White,Urban,yes,yes
+Sawyer Snyder,normal,Female,30-44,Level 4 and above,White,Rural,no,yes
+Ronny Poindexter,vegetarian,Female,60+,Level 2/Level 3/Apprenticeship/Other,BAME,Urban,yes,yes
+Messiah Lynwood,normal,Male,30-44,Level 4 and above,White,Rural,maybe,yes
+Jeremiah Millard,normal,Male,16-29,No Qualifications/Level 1,White,Urban,yes,yes
+Cherie Thorley,normal,Male,60+,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Racquel Colt,normal,Male,60+,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Ricki Courtney,normal,Female,60+,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Abe Tennyson,normal,Male,45-59,No Qualifications/Level 1,White,Urban,yes,no
+Pip Auteberry,normal,Female,30-44,Level 4 and above,BAME,Urban,yes,yes
+Dulcibella Elliott,vegetarian,Male,30-44,Level 2/Level 3/Apprenticeship/Other,White,Rural,yes,yes
+Elmer Minett,normal,Female,45-59,Level 4 and above,White,Urban,yes,yes
+Kieran Lovel,normal,Female,30-44,No Qualifications/Level 1,White,Rural,yes,yes
+Dulcibella Backus,normal,Female,16-29,Level 2/Level 3/Apprenticeship/Other,White,Urban,maybe,yes
+Lynn Madison,normal,Male,60+,No Qualifications/Level 1,White,Urban,yes,no
+Fay Thompsett,normal,Male,60+,No Qualifications/Level 1,BAME,Urban,yes,no
+Geneva Marlow,normal,Female,60+,No Qualifications/Level 1,White,Urban,yes,yes
+Savanna Revie,normal,Female,45-59,Level 2/Level 3/Apprenticeship/Other,White,Urban,maybe,yes
+Cierra Prescott,normal,Male,45-59,Level 4 and above,BAME,Urban,no,yes
+Messiah Huxley,normal,Male,30-44,Level 4 and above,White,Urban,yes,no
+Jaxson Bull,normal,Female,16-29,No Qualifications/Level 1,BAME,Rural,yes,yes
+Keegan Nye,normal,Female,45-59,Level 2/Level 3/Apprenticeship/Other,w,Rural,maybe,yes
+Janetta Foster,normal,Female,30-44,Level 2/Level 3/Apprenticeship/Other,White,Urban,no,yes
+Kian Close,normal,Male,45-59,No Qualifications/Level 1,White,Urban,maybe,no
+Trev Foss,normal,Female,30-44,No Qualifications/Level 1,White,Urban,yes,yes
+Ashleigh Whitney,normal,Female,60+,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Donald Allan,normal,Male,30-44,No Qualifications/Level 1,White,Rural,no,yes
+Hortense Platt,normal,Male,60+,No Qualifications/Level 1,White,Rural,yes,yes
+Coreen Snider,normal,Male,16-29,Level 4 and above,White,Urban,yes,yes
+Elfreda Adamson,halal,Female,45-59,Level 4 and above,White,Urban,yes,yes
+Chanelle Becket,normal,Female,30-44,Level 4 and above,White,Rural,maybe,yes
+Kilie Hunter,normal,Female,60+,Level 4 and above,White,Urban,no,yes
+Hiram Bloxam,normal,Female,60+,No Qualifications/Level 1,White,Urban,yes,yes
+Dean Jacobson,normal,Female,16-29,Level 2/Level 3/Apprenticeship/Other,White,Rural,maybe,yes
+Orpha Fairchild,normal,Female,60+,No Qualifications/Level 1,BAME,Urban,maybe,no
+Flossie Mayes,normal,Male,60+,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Ann Tennyson,normal,Female,16-29,Level 4 and above,White,Urban,yes,yes
+Suzie Appleton,vegetarian,Male,16-29,Level 4 and above,White,Urban,yes,no
+Belinda Elder,vegetarian,Female,45-59,Level 4 and above,White,Urban,yes,yes
+Page Young,normal,Female,60+,No Qualifications/Level 1,White,Urban,yes,yes
+Bertrand Whinery,normal,Male,16-29,No Qualifications/Level 1,White,Urban,maybe,no
+Arthur Blakeley,vegetarian,Female,30-44,No Qualifications/Level 1,White,Rural,yes,yes
+Mary AnnLeavitt,normal,Female,30-44,Level 2/Level 3/Apprenticeship/Other,BAME,Urban,yes,yes
+Erma Seward,normal,Female,45-59,No Qualifications/Level 1,White,Urban,yes,yes
+Gracelynn Evelyn,normal,Female,45-59,Level 4 and above,White,Urban,yes,yes
+Corine Berry,normal,Female,45-59,Level 2/Level 3/Apprenticeship/Other,White,Urban,no,yes
+Sanford Sherburn,vegan,Male,30-44,No Qualifications/Level 1,White,Rural,maybe,yes
+Kalla Carver,normal,Male,16-29,No Qualifications/Level 1,White,Rural,yes,yes
+Cleveland Hambleton,normal,Other,45-59,Level 2/Level 3/Apprenticeship/Other,White,Rural,yes,no
+Niles Jacobson,normal,Female,30-44,No Qualifications/Level 1,White,Urban,yes,yes
+Maggie Bristol,normal,Female,16-29,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Bonnie Appleby,vegetarian,Male,30-44,Level 4 and above,White,Urban,yes,yes
+Floella Drake,normal,Male,30-44,Level 2/Level 3/Apprenticeship/Other,White,Urban,maybe,yes
+Odetta Waterman,normal,Female,45-59,No Qualifications/Level 1,White,Rural,yes,yes
+Cherry Penn,normal,Male,16-29,Level 4 and above,White,Urban,no,no
+Laurel Wade,normal,Female,16-29,Level 4 and above,White,Urban,yes,yes
+Ron Moors,normal,Male,60+,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Addy Caulfield,vegetarian,Male,16-29,Level 4 and above,White,Rural,yes,yes
+Ben Cooper,normal,Female,16-29,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Eliott Woodcock,fish allergy,Female,45-59,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Amy Kingston,normal,Female,60+,No Qualifications/Level 1,White,Urban,yes,yes
+Corie Rake,normal,Female,45-59,Level 2/Level 3/Apprenticeship/Other,White,Urban,yes,yes
+Jaquelyn Reed,normal,Male,30-44,Level 2/Level 3/Apprenticeship/Other,black and minority ethnic,Urban,no,yes
+Benedict Hibbert,normal,Female,16-29,No Qualifications/Level 1,White,Urban,yes,yes