tianxuan/analysis/services/clustering.py

"""Clustering service layer: pure business logic for the clustering pipeline.

Extracted from analysis/views/clustering.py — no request/response, no rendering.
"""
import asyncio
import traceback
import logging
import warnings

import polars as pl
import numpy as np

from analysis.constants import RANDOM_SEED, UMAP_BATCH_SIZE, UMAP_TRAIN_SAMPLE, NUMERIC_DTYPES

logger = logging.getLogger(__name__)


def run_clustering_pipeline(run, store, entity_ds_id, feature_columns, algorithm,
                             min_cluster_size, run_umap=True, head=None):
    """Pure business logic — no request/response, no rendering.

    Unified clustering pipeline: clustering → feature extraction → UMAP → ORM save.

    Args:
        run: AnalysisRun ORM object.
        store: SessionStore instance.
        entity_ds_id: dataset ID in SessionStore pointing to entity-aggregated data.
        feature_columns: list of feature column names (None = auto-detect).
        algorithm: 'hdbscan' or 'kmeans'.
        min_cluster_size: int for HDBSCAN.
        run_umap: bool, whether to compute and save UMAP-2D embeddings.
        head: optional int, downsample to at most this many rows (min 100).
    """
    warnings.filterwarnings('ignore', category=RuntimeWarning, module='sklearn')

    try:
        from analysis.tool_registry import _handle_run_clustering, _handle_extract_features

        entry = store.get_dataset(entity_ds_id)
        if entry is None:
            run.status = 'failed'
            run.error_message = 'Entity dataset not found in session store'
            run.save(update_fields=['status', 'error_message'])
            return

        schema = entry.get('schema', {})

        # ── Downsampling ──
        if head is not None and head > 0:
            lf = entry['lazyframe']
            try:
                row_count = lf.select(pl.len()).collect(streaming=True).item()
            except Exception:
                logger.error("unknown failed: {}".format(traceback.format_exc()))
                row_count = 0

            if row_count > head:
                head = max(head, 100)  # ensure minimum rows for clustering
                run.progress_msg = f'正在采样 {head}/{row_count} 行...'
                run.save(update_fields=['progress_msg'])
                lf = lf.head(head)
                store.store_dataset(
                    entity_ds_id, lf,
                    schema=schema,
                    metadata=entry.get('metadata', {}),
                )
                entry = store.get_dataset(entity_ds_id)  # refresh
                row_count = head

        # Resolve feature columns — use actual LazyFrame dtypes (not stored schema dict
        # which may be stale after Utf8 coercion in _background_process).
        numeric_types = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
                         pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
                         pl.Float32, pl.Float64)

        lf = entry['lazyframe']
        live_schema = lf.collect_schema()
        live_names = live_schema.names()
        live_dtypes = list(live_schema.dtypes())

        if feature_columns:
            feature_cols = [
                c for c in feature_columns
                if c in live_names and live_dtypes[live_names.index(c)] in numeric_types
            ]
        else:
            feature_cols = [
                live_names[i] for i, dt in enumerate(live_dtypes)
                if dt in numeric_types and not live_names[i].startswith('_')
            ][:10]

        # ── Step 1: Clustering ────────────────────────────────────────────────
        # If no numeric columns detected, pass empty list to let
        # _handle_run_clustering apply its Utf8→Float64 coercion fallback.
        run.status = 'clustering'
        run.progress_pct = 70
        run.progress_msg = '正在进行聚类分析...'
        run.save(update_fields=['status', 'progress_pct', 'progress_msg'])

        result = asyncio.run(_handle_run_clustering(
            dataset_id=entity_ds_id,
            cluster_columns=feature_cols,
            algorithm=algorithm,
            params={'min_cluster_size': min_cluster_size},
            random_state=RANDOM_SEED,
        ))

        if 'error' in result:
            err = result['error']
            if 'numeric' in err.lower():
                available = [
                    f"{k}({v})" for k, v in schema.items()
                    if v.split('(')[0].strip() in numeric_types and not k.startswith('_')
                ]
                err += f"。可用数值列: {available}"
            raise Exception(err)

        cluster_id = result.get('cluster_result_id', '')
        run.cluster_count = result.get('n_clusters', 0)
        run.save(update_fields=['cluster_count'])

        # ── Step 2: Feature extraction ──────────────────────────────────────
        run.status = 'extracting'
        run.progress_pct = 90
        run.progress_msg = '正在提取聚类特征...'
        run.save(update_fields=['status', 'progress_pct', 'progress_msg'])

        feat_result = asyncio.run(_handle_extract_features(
            dataset_id=entity_ds_id,
            cluster_result_id=cluster_id,
            top_k=10, method='zscore', save_to_db=True,
            run_id=run.id,
        ))

        if 'error' in feat_result:
            raise Exception(feat_result['error'])

        # ── Step 2b: Cluster SVD feature extraction ──────────────────────
        run.progress_msg = '正在进行聚类SVD特征提取...'
        run.save(update_fields=['progress_msg'])
        try:
            from analysis.distance import cluster_svd_extract
            lf = entry['lazyframe']
            cluster_entry = store.get_cluster_result(cluster_id)
            labels_list = cluster_entry.get('labels', []) if cluster_entry else []

            if labels_list and feature_cols:
                svd_results = cluster_svd_extract(lf, labels_list, feature_cols)

                # Persist SVD features to ClusterFeature (method='cluster_svd')
                from analysis.models import ClusterFeature as CF
                from analysis.models import ClusterResult as CR
                for label, svd_info in svd_results.items():
                    cr = CR.objects.filter(run=run, cluster_label=label).first()
                    if cr is None:
                        continue
                    # Delete old zscore entries for this cluster, replace with SVD
                    CF.objects.filter(cluster=cr, distinguishing_method='zscore').delete()
                    for feat_name in svd_info.get('features', [])[:10]:
                        idx = svd_info['features'].index(feat_name) if feat_name in svd_info['features'] else -1
                        strength = (svd_info.get('feature_strength', [])[idx]
                                    if idx >= 0 and idx < len(svd_info.get('feature_strength', []))
                                    else None)
                        CF.objects.get_or_create(
                            cluster=cr,
                            feature_name=feat_name,
                            defaults={
                                'mean': None, 'std': None, 'median': None,
                                'p25': None, 'p75': None, 'missing_rate': None,
                                'distinguishing_score': float(strength) if strength is not None else None,
                                'distinguishing_method': 'cluster_svd',
                            },
                        )
        except Exception as svd_err:
            logger.warning(f'Cluster SVD extraction skipped: {svd_err}')

        # ── Step 3: UMAP-3D embedding (2D fallback) ───────────────────────
        if run_umap:
            try:
                lf = entry['lazyframe']
                # UMAP: sample max 10K for training, batch-transform rest in 1K batches
                MAX_UMAP_TRAIN = UMAP_TRAIN_SAMPLE
                BATCH_SIZE = UMAP_BATCH_SIZE
                try:
                    n_total = lf.select(pl.len()).collect(streaming=True).item()
                except Exception:
                    logger.error("unknown failed: {}".format(traceback.format_exc()))
                    n_total = 0
                from sklearn.preprocessing import StandardScaler
                import umap

                live_schema = lf.collect_schema()
                num_cols = [name for name, dt in zip(live_schema.names(), live_schema.dtypes())
                            if dt in (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
                                      pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
                                      pl.Float32, pl.Float64) and not name.startswith('_')]

                if len(num_cols) >= 2:
                    if n_total > MAX_UMAP_TRAIN:
                        run.progress_msg = f'UMAP 采样 {MAX_UMAP_TRAIN}/{n_total} 实体训练...'
                        run.save(update_fields=['progress_msg'])
                        df_sample = lf.sample(n=MAX_UMAP_TRAIN, seed=RANDOM_SEED).collect(streaming=True)
                        df_umap = lf.collect(streaming=True)
                    else:
                        df_umap = lf.collect(streaming=True)

                    coords, umap_components = compute_umap_embedding(
                        df_umap, num_cols, n_total,
                        max_train=MAX_UMAP_TRAIN,
                        batch_size=BATCH_SIZE,
                        seed=RANDOM_SEED,
                    )

                    cluster_entry = store.get_cluster_result(cluster_id)
                    labels_list = cluster_entry.get('labels', []) if cluster_entry else []

                    save_entity_profiles(run, df_umap, labels_list, coords, umap_components)
            except Exception as umap_err:
                logger.warning(f'UMAP embedding skipped: {umap_err}')

        # Fallback: if entity_count not set (UMAP skipped), derive from cluster sizes
        if run.entity_count is None:
            from django.db.models import Sum
            total = run.clusters.aggregate(total=Sum('size'))['total'] or 0
            run.entity_count = total
            run.save(update_fields=['entity_count'])

        # ── Done ──
        run.status = 'completed'
        run.progress_pct = 100
        run.progress_msg = '分析完成'
        run.save(update_fields=['status', 'progress_pct', 'progress_msg', 'entity_count'])

    except Exception:
        tb = traceback.format_exc()
        logger.error(tb)
        try:
            run.status = 'failed'
            run.error_message = tb
            run.progress_msg = f'失败: {tb}'
            run.run_log += f'\n[ERROR] {tb}'
            run.save(update_fields=['status', 'error_message', 'progress_msg', 'run_log'])
        except Exception as e:
            logger.warning('save failed after pipeline error: %s', e)


def compute_umap_embedding(df, num_cols, n_total, max_train, batch_size, seed):
    """UMAP 3D→2D fallback, returns (coords, umap_components).

    Args:
        df: Collected Polars DataFrame with entity-level data.
        num_cols: list of numeric column names to use for UMAP.
        n_total: total row count.
        max_train: max rows for UMAP training sample.
        batch_size: batch size for transform when dataset exceeds max_train.
        seed: random seed.

    Returns:
        (coords, umap_components): numpy array of shape (n_total, umap_components),
        and umap_components is 2 or 3.
    """
    from sklearn.preprocessing import StandardScaler
    import umap

    if len(df) <= 2 or len(num_cols) < 2:
        return None, 0

    if n_total > max_train:
        df_sample = df.sample(n=max_train, seed=seed)
    else:
        df_sample = df

    umap_components = 3  # try 3D first
    coords = None

    for attempt_n in (3, 2):
        try:
            if n_total > max_train:
                # Train UMAP on sample, batch-transform full dataset
                mat_sample = df_sample.select(num_cols).to_numpy()
                scaler = StandardScaler().fit(mat_sample)
                mat_sample_scaled = np.nan_to_num(
                    scaler.transform(mat_sample), nan=0.0)
                reducer = umap.UMAP(n_components=attempt_n,
                                    random_state=seed,
                                    n_neighbors=15, min_dist=0.1,
                                    metric='euclidean')
                reducer.fit(mat_sample_scaled)
                # Batch transform in batch_size-row batches
                coords_list = []
                for start in range(0, len(df), batch_size):
                    end = min(start + batch_size, len(df))
                    mat_batch = scaler.transform(
                        df[start:end].select(num_cols).to_numpy())
                    mat_batch = np.nan_to_num(mat_batch, nan=0.0)
                    coords_list.append(reducer.transform(mat_batch))
                coords = np.vstack(coords_list)
            else:
                mat = np.nan_to_num(StandardScaler().fit_transform(
                    df.select(num_cols).to_numpy()), nan=0.0)
                reducer = umap.UMAP(n_components=attempt_n,
                                    random_state=seed,
                                    n_neighbors=15, min_dist=0.1,
                                    metric='euclidean')
                coords = reducer.fit_transform(mat)
            umap_components = attempt_n
            break
        except Exception as e_umap:
            if attempt_n == 2:
                raise
            logger.warning(f'UMAP 3D failed ({e_umap}), falling back to 2D')

    if coords is None:
        raise Exception('UMAP produced no output')

    return coords, umap_components


def save_entity_profiles(run, df_umap, labels_list, coords, umap_components):
    """Persist per-row cluster labels + UMAP coords to EntityProfile.

    Args:
        run: AnalysisRun ORM object.
        df_umap: Collected Polars DataFrame (used for entity column detection + row iteration).
        labels_list: list of cluster labels per row.
        coords: numpy array of UMAP coordinates (must not be None).
        umap_components: 2 or 3.
    """
    from analysis.models import EntityProfile as EP
    from analysis.models import ClusterResult as CR

    # Determine entity column (first non-numeric, non-underscore column)
    df_num_cols = [c for c, dt in zip(df_umap.columns, df_umap.dtypes)
                   if dt in NUMERIC_DTYPES]
    non_num = [c for c in df_umap.columns
               if c not in df_num_cols and not c.startswith('_')]
    ent_col = non_num[0] if non_num else df_umap.columns[0]

    profiles = []
    cr_cache = {}
    for i, row in enumerate(df_umap.iter_rows(named=True)):
        ev = str(row.get(ent_col, '')) or 'entity'
        ev = f'{ev}_{i}'
        lbl = labels_list[i] if i < len(labels_list) else -1
        cache_key = f'{run.id}_{lbl}'
        if cache_key not in cr_cache:
            cr_cache[cache_key] = CR.objects.filter(
                run=run, cluster_label=lbl).first()
        has_z = umap_components >= 3 and coords.shape[1] >= 3
        profiles.append(EP(
            entity_value=ev, run=run, cluster_label=lbl,
            cluster=cr_cache[cache_key],
            embedding_x=float(coords[i, 0]) if i < len(coords) else None,
            embedding_y=float(coords[i, 1]) if i < len(coords) else None,
            embedding_z=float(coords[i, 2]) if has_z
                and i < len(coords) else None,
        ))
    if profiles:
        EP.objects.bulk_create(profiles, ignore_conflicts=True, batch_size=1000)
        run.entity_count = len(profiles)
        run.save(update_fields=['entity_count'])