diffnarm_wrapper.py

"""
This script allows you to run DiffNARM on your own dataset.

Run DiffNARM on your data from the command line as follows:
    python run_diffnarm.py path/to/your_data.csv
    python run_diffnarm.py path/to/your_data.csv --output rules.json

DiffNARM can also imported into your Python ML pipeline as follows (requires diffnarm and utils packages in this repo)
    import pandas as pd
    from run_diffnarm import mine_rules

    result = mine_rules(pd.read_csv("your_data.csv"))
    result["rules"]        # list of rules with quality metrics
    result["statistics"]   # aggregate statistics (rule count, coverage, averages)

The input is a plain table (CSV or DataFrame), one row per transaction. Numerical
and categorical columns are detected automatically: numerical columns are binned
with learned boundaries, categorical columns are one-hot encoded.
"""
import argparse
import json

import numpy as np
import pandas as pd

from src.diffnarm.model import train
from src.diffnarm.rule_extraction import generate_rules

DEFAULT_PARAMS = {
    # ── discretization ─────────────────────────────────────────────────────────
    'n_bins': 5,  # bins per numerical feature; also accepts a dict {column_name: bins}
    # initial boundary placement before training: 'equal_freq', 'equal_width', 'kmeans', 'zscore', 'jenks', 'gmm', 'optbinning'
    'init_discretizer': 'equal_width',

    # ── training ───────────────────────────────────────────────────────────────
    'epochs': 5,  # training epochs
    'lr': 5e-3,  # learning rate
    'mask_prob': 0.75,  # probability of masking each feature per training step;
    'lambda_bce': 1,  # weight of the reconstruction (BCE) loss
    'lambda_reg': 10,  # weight of the regression loss on numerical values;

    # weight of the confidence reconstruction loss. This is disabled by default, as we found that it is less effective than a simple regression loss
    'lambda_conf_recon': 0,

    'beta_start': 1.0,  # binning sharpness at the first epoch (soft bin memberships)
    'beta_end': 10.0,  # sharpness at the last epoch; annealed geometrically in between,

    # where the regression loss reads a feature's value from: 'own_bins_only' (the feature's own bin activations, default),
    #   'encoder' or 'decoder'. Only use 'own_bins_only' as we found the other options to be underperforming.
    'reg_placement': 'own_bins_only',

    # whether the reconstruction loss should update boundaries, or should it be regression loss only
    'ae_updates_boundaries': True,

    'device': None,  # None = use GPU if available; or 'cpu' / 'cuda'

    # ── rule extraction ────────────────────────────────────────────────────────
    'max_antecedents': 2,  # maximum antecedent length; None = grow until no
    # antecedent passes the frequency gate
    'min_estimates': {
        'rule_frequency': 0.25,  # minimum estimated frequency of the antecedent
        'rule_strength': 0.8,  # minimum implication probability of the consequent
    },

    # estimate antecedent frequency from pairwise implication probabilities (geometric mean) instead of the min single
    #   feature value frequency
    'use_pairwise_support': True,
}


def _find_collapsed_features(model):
    """Features whose learned bins collapsed (effective bin count below half nominal)."""
    cont_cols = {col for info in model._groups.values() for col in info['cols']}
    hard = model.input_vectors.to_numpy()
    collapsed = set()
    for fvi in model.feature_value_indices:
        if fvi['feature'] not in cont_cols:
            continue
        k = fvi['end'] - fvi['start']
        bin_occ = hard[:, fvi['start']:fvi['end']].sum(axis=0)
        bin_occ = bin_occ / (bin_occ.sum() + 1e-10)
        entropy = float(-np.sum(bin_occ * np.log(bin_occ + 1e-10)))
        if np.exp(entropy) < k / 2:
            collapsed.add(fvi['feature'])
    return collapsed


def mine_rules(data, **overrides):
    """Train DiffNARM on a dataset and extract association rules.

    :param data: pandas DataFrame or path to a CSV file
    :param overrides: any key from DEFAULT_PARAMS, e.g. mine_rules(df, n_bins=3, epochs=10)
    :return: dict with 'rules' and 'statistics'
    """
    unknown = set(overrides) - set(DEFAULT_PARAMS)
    if unknown:
        raise ValueError(f"Unknown parameter(s): {sorted(unknown)}. See DEFAULT_PARAMS in {__file__}")
    params = {**DEFAULT_PARAMS, **overrides}

    df = pd.read_csv(data) if not isinstance(data, pd.DataFrame) else data

    model = train(
        df,
        n_bins=params['n_bins'],
        init_discretizer=params['init_discretizer'],
        epochs=params['epochs'],
        lr=params['lr'],
        mask_prob=params['mask_prob'],
        lambda_bce=params['lambda_bce'],
        lambda_reg=params['lambda_reg'],
        lambda_conf_recon=params['lambda_conf_recon'],
        beta_start=params['beta_start'],
        beta_end=params['beta_end'],
        reg_placement=params['reg_placement'],
        ae_updates_boundaries=params['ae_updates_boundaries'],
        device=params['device'],
    )
    if model is None:
        return {'rules': [], 'statistics': {}}

    # exclude features whose bins collapsed during training from both rule sides
    usable = [fvi['feature'] for fvi in model.feature_value_indices
              if fvi['feature'] not in _find_collapsed_features(model)]

    return generate_rules(
        model,
        max_antecedents=params['max_antecedents'],
        min_estimates=params['min_estimates'],
        use_pairwise_support=params['use_pairwise_support'],
        features_of_interest=usable,
        target_classes=usable,
    )


def _format_rule(rule):
    ant = " AND ".join(f"{a['feature']} = {a['value']}" for a in rule['antecedents'])
    cons = f"{rule['consequent']['feature']} = {rule['consequent']['value']}"
    metrics = ", ".join(f"{m}={rule[m]:.3f}" for m in ('support', 'confidence', 'zhangs_metric')
                        if m in rule)
    return f"IF {ant} THEN {cons}  ({metrics})"


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Run DiffNARM on a CSV dataset.")
    parser.add_argument('csv', help="path to the input CSV file (one row per transaction)")
    parser.add_argument('--output', help="write the rules and statistics to this JSON file")
    parser.add_argument('--n-bins', type=int, default=DEFAULT_PARAMS['n_bins'],
                        help="bins per numerical feature")
    parser.add_argument('--epochs', type=int, default=DEFAULT_PARAMS['epochs'],
                        help="training epochs")
    parser.add_argument('--max-antecedents', type=int, default=DEFAULT_PARAMS['max_antecedents'],
                        help="maximum antecedent length")
    args = parser.parse_args()

    result = mine_rules(args.csv, n_bins=args.n_bins, epochs=args.epochs,
                        max_antecedents=args.max_antecedents)

    stats = result['statistics']
    print(f"\n{len(result['rules'])} rules"
          + (f", data coverage={stats['data_coverage']:.3f}" if 'data_coverage' in stats else ""))
    for key, value in stats.items():
        if key.startswith('average_'):
            print(f"  {key}: {value:.3f}")
    print()
    for rule in result['rules']:
        print(_format_rule(rule))

    if args.output:
        with open(args.output, 'w') as f:
            json.dump(result, f, indent=2, default=float)
        print(f"\nSaved to {args.output}")