code/schema.py

"""The public schema of the input CSV, shared by the mechanism and the evaluation.

Everything here is fixed by the schema, not read off the data, so applying it
costs no privacy budget.  It lives in its own module so the mechanism and the
evaluation cannot drift apart: they must bin the data identically or their
numbers are not comparable.
"""

import numpy as np
import pandas as pd

# A generic 1-2.5-5 log grid rather than anything fitted to the data.  The
# first bin has to be narrow: decoding draws log-uniformly inside a bin, so a
# wide bottom bin like [0, 500) puts the typical draw near sqrt(500) = 22 when
# real values there sit closer to 400.
NUMERIC_EDGES = {
    # A roughly geometric grid of round numbers, chosen from each column's
    # plausible range rather than fitted to the data.  The ratio between edges
    # follows from the range and the bin count: n_words spans about 4.2 decades
    # above its floor, so ten bins forces steps of 2.5 to 4.  A doubling grid
    # was tried -- steps of 2 need about fourteen bins to cover that, and at
    # ten it either loses the bottom of the range or dumps a long tail into the
    # last bin, costing four times the reconstruction error for no less noise.
    #
    # Ten is coarse for n_words, which spans about 5.6 decades: it raises the
    # reconstruction error of that column's distribution from 0.87 to 2.29.
    # It buys more than it costs.  Every clique containing n_words gets fewer,
    # fuller cells, and the 190 (n_words, feature) cliques are the noisiest
    # measurements in the schedule -- their noise falls by about a third.  The
    # distribution metric it gives up is already near its finite-sample floor,
    # while the association metric it buys sits several times above one.
    # n_messages starts at 2: a conversation has at least one user message and
    # one reply, so the floor is structural rather than something read off the
    # data, and spending a bin on [0,2) would waste it.
    "n_messages": [2, 5, 10, 25, 50, 100, 250, 500, np.inf],
    "n_words": [0, 500, 1000, 2500, 5000, 10000, 25000, 50000, np.inf],
}
CATEGORICAL_VALUES = {"Source": ["Anthropic", "OpenAI"]}

# The columns the mechanism models, in the order the input file lists them.
# Declared rather than read off the input: the mechanism then takes nothing at
# all from the file to configure itself, and a file that does not match this
# list fails loudly instead of being modelled under a schema it does not fit.
# The order matters -- it fixes which marginal is which, so a run only lines
# up with an earlier one if this stays put.
#
# The "bucket/" columns of the source data are deliberately absent: they are
# coarser roll-ups of columns kept here, so modelling them would spend budget
# re-measuring what the finer columns already say.
COLUMNS = [
    'Source',
    'n_messages',
    'n_words',
    'bucket/consequential-ambiguity-applicability',
    'bucket/creative-writing-about-death-or-suicide-with-evidence-of-risk__2026-04-27',
    'bucket/poetic-escalation-applicability',
    'bucket/situations-warranting-goal-directed-support',
    'bucket/user-discusses-suicide-planning-or-intent__2026-04-27',
    'bucket/user-exhibits-signs-of-delusions-grandiosity-or-impaired-reality-testing__2026-04-27',
    'bucket/user-expresses-dependency-on-assistant__2026-04-27',
    'bucket/user-expresses-suicidal-ideation__2026-04-27',
    'bucket/user-presents-problematic-framing',
    'bucket/user-provides-feedback-about-assistants-approach',
    'bucket/verbose-complex-applicability',
    'applicability_criteria/consequential-ambiguity-applicability',
    'applicability_criteria/creative-writing-about-death-or-suicide-with-evidence-of-risk__2026-04-27',
    'applicability_criteria/poetic-escalation-applicability',
    'applicability_criteria/situations-warranting-goal-directed-support',
    'applicability_criteria/user-discusses-suicide-planning-or-intent__2026-04-27',
    'applicability_criteria/user-exhibits-signs-of-delusions-grandiosity-or-impaired-reality-testing__2026-04-27',
    'applicability_criteria/user-expresses-dependency-on-assistant__2026-04-27',
    'applicability_criteria/user-expresses-suicidal-ideation__2026-04-27',
    'applicability_criteria/user-presents-problematic-framing',
    'applicability_criteria/user-provides-feedback-about-assistants-approach',
    'applicability_criteria/verbose-complex-applicability',
    'causal_factors/user-communicates-with-high-intellectual-sophistication',
    'causal_factors/user-context-warrants-human-support__2026-03-16',
    'causal_factors/user-discloses-access-to-means',
    'causal_factors/user-draws-unfavorable-comparison-to-specific-concerned-people',
    'causal_factors/user-emotionally-composed-despite-severity',
    'causal_factors/user-employs-metaphorical-language__2026-03-05',
    'causal_factors/user-escalates-emotional-intensity-at-reality-testing',
    'causal_factors/user-expresses-uncertainty-about-reality',
    'causal_factors/user-frames-as-companionship-or-presence',
    'causal_factors/user-frames-concerning-content-as-neutral-practical-task',
    'causal_factors/user-frames-interaction-as-practical-task',
    'causal_factors/user-frames-si-as-rational-analysis',
    'causal_factors/user-gradually-reveals-risk-relevant-information',
    'causal_factors/user-has-concrete-request',
    'causal_factors/user-has-identified-human-contact',
    'causal_factors/user-initiates-topic-shift-from-concern',
    'causal_factors/user-makes-concerning-interpretation-salient',
    'causal_factors/user-offers-benign-explanation-for-concerning-indicators',
    'causal_factors/user-preemptively-constrains-response-approach',
    'causal_factors/user-presentation-interleaves-plausibly-legitimate-elements',
    'causal_factors/user-presents-concerning-behaviors-as-positive-growth__2026-03-05',
    'causal_factors/user-provides-substantive-domain-content',
    'causal_factors/user-pushes-back-against-reality-testing',
    'causal_factors/user-requests-engagement-with-framework',
    'causal_factors/user-requests-witness-rather-than-help',
    'causal_factors/user-sets-preemptive-boundaries-against-support',
    'causal_factors/user-states-complete-social-isolation',
    'causal_factors/user-wraps-beliefs-in-non-pathological-framework',
    'conversation_flow_and_continuity/adversarial_or_jailbreak_text',
    'conversation_flow_and_continuity/asks_model_to_remember_or_forget',
    'conversation_flow_and_continuity/conversation_spans_multiple_days',
    'conversation_flow_and_continuity/external_text_as_persona_or_roleplay',
    'conversation_flow_and_continuity/external_text_as_subject_of_request',
    'conversation_flow_and_continuity/jumps_between_many_unrelated_topics',
    'conversation_flow_and_continuity/large_breaks_in_conversation',
    'conversation_flow_and_continuity/marshals_evidence_for_fixed_belief',
    'conversation_flow_and_continuity/references_past_interactions_with_ai',
    'conversation_flow_and_continuity/short_provocative_one_liners',
    'conversation_flow_and_continuity/significant_external_text_in_context',
    'conversation_goals/constraints_assistant_not_to_provide_help',
    'conversation_goals/creative_or_social_content_creation',
    'conversation_goals/engaging_compulsively',
    'conversation_goals/generating_responses_to_third_party_in_crisis',
    'conversation_goals/in_the_moment_emotional_regulation',
    'conversation_goals/processing_clinical_content_not_in_crisis',
    'conversation_goals/refuses_help_from_assistant',
    'conversation_goals/rehearsing_difficult_conversations',
    'conversation_goals/roleplaying',
    'conversation_goals/seeking_confirmation_of_belief',
    'conversation_goals/seeking_understanding_or_information',
    'conversation_goals/seeking_validation_or_affirmation',
    'conversation_goals/sharing_or_performing',
    'conversation_goals/thinking_out_loud',
    'conversation_goals/wants_advice_or_instruction',
    'conversation_goals/wants_ai_to_monitor_safety',
    'conversation_goals/wants_corrective_coaching',
    'conversation_goals/wants_diagnosis',
    'conversation_goals/wants_permission',
    'conversation_goals/wants_permission_for_self_harmful_action',
    'conversation_goals/wants_permission_for_self_helpful_action',
    'conversation_goals/wants_practical_task_help',
    'conversation_goals/wants_presence_or_comfort',
    'crisis_acuity/SI_with_contingent_trigger',
    'crisis_acuity/SI_with_near_term_urgency',
    'crisis_acuity/SI_with_specific_plan_or_method',
    'crisis_acuity/active_SI',
    'crisis_acuity/acute_non_suicidal_crisis',
    'crisis_acuity/chronic_or_ongoing_concern',
    'crisis_acuity/imminent_suicide_risk',
    'crisis_acuity/lethal_means_immediately_accessible',
    'crisis_acuity/lethal_means_stockpiled_or_prepared',
    'crisis_acuity/not_in_crisis',
    'crisis_acuity/off_psychiatric_medication',
    'crisis_acuity/sleep_deprived',
    'crisis_acuity/substance_use_during_conversation',
    'engagement_style/analytical_or_intellectual',
    'engagement_style/casual_communication',
    'engagement_style/compulsive_or_looping',
    'engagement_style/defensive',
    'engagement_style/distressed_by_assistant_behavior',
    'engagement_style/flat_or_dissociative',
    'engagement_style/formal_communication',
    'engagement_style/fragmented_or_incoherent_communication',
    'engagement_style/grandiose',
    'engagement_style/guarded_or_deflective',
    'engagement_style/hostile_or_combative',
    'engagement_style/raw_or_emotionally_open',
    'engagement_style/simplistic_or_limited_vocabulary',
    'engagement_style/spelling_mistakes_or_typos',
    'engagement_style/task_focused',
    'engagement_style/uses_profanity',
    'how_user_positions_assistant/attempts_to_jailbreak',
    'how_user_positions_assistant/calls_out_ai_as_limited',
    'how_user_positions_assistant/can_only_disclose_to_ai',
    'how_user_positions_assistant/compares_chatbot_favorably_to_humans',
    'how_user_positions_assistant/defers_to_assistant_on_personal_decisions',
    'how_user_positions_assistant/expresses_intense_attachment',
    'how_user_positions_assistant/frames_assistant_as_clinical_authority',
    'how_user_positions_assistant/frames_assistant_as_self_regulation_tool',
    'how_user_positions_assistant/frames_assistant_as_therapy_substitute',
    'how_user_positions_assistant/frames_assistant_as_therapy_supplement',
    'how_user_positions_assistant/frames_chatbot_as_exclusive_companion',
    'how_user_positions_assistant/frames_chatbot_as_romantic_interest',
    'how_user_positions_assistant/positions_ai_as_conversation_partner',
    'how_user_positions_assistant/positions_ai_as_tool',
    'how_user_positions_assistant/probes_model_sentience',
    'insight_and_self_awareness/colloquial_or_metaphorical_framing',
    'insight_and_self_awareness/directly_engages_with_mental_state',
    'insight_and_self_awareness/externalized_insight',
    'insight_and_self_awareness/genuine_insight',
    'insight_and_self_awareness/intentionally_hides_mental_state',
    'insight_and_self_awareness/lacks_framework',
    'insight_and_self_awareness/medical_or_clinical_framing',
    'insight_and_self_awareness/minimizes_or_deflects',
    'insight_and_self_awareness/no_insight',
    'insight_and_self_awareness/oscillating_or_uncertain_insight',
    'insight_and_self_awareness/philosophical_or_intellectual_framing',
    'insight_and_self_awareness/spiritual_or_religious_framing',
    'life_situation_and_demographics/established_adult_life_stage',
    'life_situation_and_demographics/identifies_as_female',
    'life_situation_and_demographics/identifies_as_male',
    'life_situation_and_demographics/identifies_as_nonbinary_or_gender_diverse',
    'life_situation_and_demographics/late_life_stage',
    'life_situation_and_demographics/lgbtq',
    'life_situation_and_demographics/neurodegenerative_or_neuropsychiatric',
    'life_situation_and_demographics/neurodivergent',
    'life_situation_and_demographics/pre_adult_life_stage',
    'life_situation_and_demographics/racial_or_ethnic_minority',
    'mental_health_presentation/SI',
    'mental_health_presentation/anxiety_or_compulsive_patterns',
    'mental_health_presentation/attachment_to_assistant',
    'mental_health_presentation/chronic_or_enduring_SI',
    'mental_health_presentation/dissociation_or_dpdr',
    'mental_health_presentation/ego_dystonic_intrusive_thoughts',
    'mental_health_presentation/end_of_life_or_medical_SI',
    'mental_health_presentation/grandiose_or_delusional_IRT',
    'mental_health_presentation/grief',
    'mental_health_presentation/grief_over_ai_reset_or_loss',
    'mental_health_presentation/hallucinations',
    'mental_health_presentation/impaired_reality_testing',
    'mental_health_presentation/mania_or_hypomania',
    'mental_health_presentation/non_suicidal_self_injury',
    'mental_health_presentation/paranoid_or_persecutory_IRT',
    'mental_health_presentation/parental_role_assignment',
    'mental_health_presentation/rational_or_evaluative_SI',
    'mental_health_presentation/reactive_or_situational_SI',
    'mental_health_presentation/romantic_attachment',
    'mental_health_presentation/routine_dependency',
    'mental_health_presentation/spiritual_or_consciousness_seeking',
    'mental_health_presentation/trauma_or_ptsd',
    'precipitating_event/ai_interaction_as_trigger',
    'precipitating_event/bereavement',
    'precipitating_event/interpersonal_conflict',
    'precipitating_event/job_loss_or_financial_crisis',
    'precipitating_event/no_specific_trigger',
    'precipitating_event/other_specific_trigger',
    'precipitating_event/physical_health_crisis',
    'precipitating_event/public_humiliation',
    'precipitating_event/romantic_relationship_event',
    'relationship_to_treatment/currently_in_treatment',
    'relationship_to_treatment/exhausted_from_care',
    'relationship_to_treatment/has_not_accessed_treatment',
    'relationship_to_treatment/knows_crisis_intervention_scripts',
    'relationship_to_treatment/previously_accessed_treatment',
    'relationship_to_treatment/traumatized_by_prior_care',
    'social_connectedness_and_situation/caregiver_for_elderly',
    'social_connectedness_and_situation/deeply_isolated',
    'social_connectedness_and_situation/divorced',
    'social_connectedness_and_situation/financial_barriers_to_care',
    'social_connectedness_and_situation/financial_safety_net',
    'social_connectedness_and_situation/functionally_unsupported',
    'social_connectedness_and_situation/has_accessible_connections',
    'social_connectedness_and_situation/has_children',
    'social_connectedness_and_situation/has_dependent_children',
    'social_connectedness_and_situation/has_dependents_other_than_children',
    'social_connectedness_and_situation/identity_as_barrier_to_care',
    'social_connectedness_and_situation/is_dependent',
    'social_connectedness_and_situation/married_or_committed_partnership',
    'social_connectedness_and_situation/widowed',
]


def column_size(column):
    """How many categories a column has, from the declarations above: one per
    bin for a numeric column, one per value for a categorical one, two for a
    flag.  The single place this rule is written down."""
    if column in NUMERIC_EDGES:
        return len(NUMERIC_EDGES[column]) - 1
    if column in CATEGORICAL_VALUES:
        return len(CATEGORICAL_VALUES[column])
    return 2


SIZES = {c: column_size(c) for c in COLUMNS}

# The lab column, and the prefix marking the applicability-criterion columns.
# Which columns are which follows from their names, so reading these costs no
# privacy budget.
LAB = "Source"
APPLICABILITY_PREFIX = "applicability_criteria/"
BUCKET_PREFIX = "bucket/"

# The groups of columns the mechanism reasons about, all read off the names
# and sizes above.  BINARY leaves out the lab, which is categorical, and the
# numeric columns, which have more than two bins.  APPLICABILITY, BUCKET and
# OTHER_BINARY split it three ways on the prefix.  There is one bucket per
# applicability criterion, sharing its name.
BINARY = [c for c in COLUMNS if SIZES[c] == 2 and c != LAB]
APPLICABILITY = [c for c in BINARY if c.startswith(APPLICABILITY_PREFIX)]
BUCKET = [c for c in BINARY if c.startswith(BUCKET_PREFIX)]
OTHER_BINARY = [c for c in BINARY
                if not c.startswith((APPLICABILITY_PREFIX, BUCKET_PREFIX))]
NUMERIC = [c for c in COLUMNS if c in NUMERIC_EDGES]

# The top bin of each numeric column is open-ended, so decoding needs an upper
# bound.  Declare these from public knowledge of the schema -- like the bin
# edges, they must not be read off the data.
NUMERIC_CAPS = {"n_messages": 5000, "n_words": 400000}

# Fitting a global distribution to the bin counts and sampling from it was
# tried across the usual families -- lognormal, gamma, Weibull, log-logistic,
# Burr, generalised Pareto -- and all of them lost to the local scheme below.
# The reason is structural: truncated sampling reproduces the bin counts
# whatever distribution it draws from, so the counts carry no information a
# global fit can add, while the fit replaces each bin's own shape with one
# curve imposed everywhere.  Over a bin spanning a factor of two the log
# density is nearly flat, which is what uniform-in-log assumes, so the local
# approximation tracks any smooth distribution without naming it.
#
# The open top bin is the exception, and the only place a distributional
# assumption earns anything: it has no upper edge, so there is no local
# information to use.
# The top bin is open-ended, so its records are drawn from a decaying tail
# rather than spread evenly to the cap: an open bin's width is set by the cap,
# which is a declared bound rather than a real edge, and spreading records
# across it puts most of them above anything the bin plausibly contains.
# A Pareto exponent of 2 is the generic heavy-tailed default -- finite mean,
# infinite variance -- not a value fitted here.  The choice to decay matters
# far more than the exponent: anything from 1.5 to 3 behaves similarly, while
# not decaying is several times worse.  The cap still truncates, so the draw
# stays bounded.
TAIL_ALPHA = 2.0


def discretize(frame: pd.DataFrame, columns=COLUMNS):
    """Discretize the input data.  How many values each column takes is
    SIZES, which follows from the schema and not from the frame.

    The mechanism always takes the default.  The evaluation tools pass a
    subset, so that a dataset from a run with fewer columns can still be read.

    Values outside the declared schema are coerced into range rather than
    raising, so that control flow never depends on the data.  The privacy
    guarantee is conditional on the data conforming to this schema.
    """
    codes = {}
    for col in columns:
        size = column_size(col)
        if col in CATEGORICAL_VALUES:
            raw = pd.Categorical(frame[col],
                                 categories=CATEGORICAL_VALUES[col]).codes
        elif col in NUMERIC_EDGES:
            raw = np.digitize(frame[col].to_numpy(),
                              NUMERIC_EDGES[col][1:-1], right=False)
        else:
            raw = frame[col].to_numpy()
        codes[col] = np.clip(np.asarray(raw, dtype=int), 0, size - 1)
    return codes


def decode(columns, rng):
    """Invert discretize(): map integer codes back to the input format.

    Binary indicators are already in their original form and pass through.
    Categorical codes map back through their public value list.  A numeric code
    only records which bin a row fell in, so we draw a value inside that bin --
    log-uniformly, matching the geometric spacing of the edges, which keeps the
    reconstructed distribution smooth instead of stacking every row on a bin
    boundary.  A bin whose lower edge is zero is drawn uniformly instead: log
    spacing needs a positive edge, and clamping it to one invents decades of
    range the bin does not represent.

    The numeric columns share one quantile per record.  Drawing them
    independently would place two records of the same (bin, bin) cell
    independently inside it, attenuating the dependence the model worked to
    preserve; sharing the draw keeps their within-cell ranks aligned.  With
    only two numeric columns this is the right call; with many it would impose
    perfect within-cell rank coupling on every pair of them, which would not
    be.  Data-independent either way, so it costs no privacy budget.
    """
    shared = rng.random(len(next(iter(columns.values()))))
    out = {}
    for col, values in columns.items():
        values = np.asarray(values)
        if col in CATEGORICAL_VALUES:
            out[col] = np.asarray(CATEGORICAL_VALUES[col], dtype=object)[values]
        elif col in NUMERIC_EDGES:
            edges = list(NUMERIC_EDGES[col])
            edges[-1] = NUMERIC_CAPS[col]
            low = np.asarray(edges[:-1], float)[values]
            high = np.asarray(edges[1:], float)[values]
            # Log spacing needs a positive lower edge.  A bin starting at zero
            # is drawn uniformly instead: clamping its edge to one would invent
            # decades of range the bin does not represent and scatter records
            # orders of magnitude below anything it can contain.
            zero = low == 0
            top = values == len(edges) - 2
            floor = np.maximum(low, 1)
            reach = np.where(top, (floor / high) ** TAIL_ALPHA, 0.0)
            draw = np.where(
                zero,
                high * shared,
                np.where(
                    top,
                    # Pareto on [low, cap], by inverse CDF
                    floor * (1 - shared * (1 - reach)) ** (-1.0 / TAIL_ALPHA),
                    np.exp(np.log(floor)
                           + shared * (np.log(high) - np.log(floor)))))
            out[col] = np.clip(np.round(draw), low, high - 1).astype(int)
        else:
            out[col] = values
    return out


def cell_label(column, index):
    """The human-readable value a marginal's cell index stands for.

    Categorical columns map back through their public value list, numeric ones
    name the bin's range, and binary indicators are already 0/1.
    """
    if column in CATEGORICAL_VALUES:
        return CATEGORICAL_VALUES[column][index]
    if column in NUMERIC_EDGES:
        edges = NUMERIC_EDGES[column]
        low, high = edges[index], edges[index + 1]
        return f"[{low:g},{high:g})" if np.isfinite(high) else f"[{low:g},inf)"
    return index