mtam.py

"""
MTAM: Manuscript–Trope Affiliation Model.
Provides SBM fitting and management over bipartite chant graphs.
"""

import graph_tool.all as gt
import numpy as np
import matplotlib.pyplot as plt
import pickle
import time


class mtam:
    """Manuscript–Trope Affiliation Model with SBM support."""

    def __init__(self, *, deg_corr=True, nested=True, anneal=True,
                 max_zero_sweeps=100000, plateau_patience=5000,
                 n_threads=8, omp_schedule="dynamic"):
        self.g = None
        self.trope_elements = []
        self.manuscripts = []
        self.states = []
        self.groups = {}
        self.mdl = np.nan
        self.mdl_history = []

        self.deg_corr = deg_corr
        # Not in this version
        self.nested = nested
        self.anneal = anneal
        self.max_zero_sweeps = max_zero_sweeps
        self.plateau_patience = plateau_patience
        self.n_threads = n_threads
        self.omp_schedule = omp_schedule

    def load_network(self, path):
        """Load graph-tool graph."""
        self.g = gt.load_graph(path)
        print(f"Loaded graph with {self.g.num_vertices()} vertices, {self.g.num_edges()} edges")

    def fit(self, n_init=10, model_type="SBM"):
        """Fit SBM multiple times and keep best result."""
        self.states = []
        self.mdl_history = []

        for i in range(n_init):
            args = dict(deg_corr=self.deg_corr, clabel=self.g.vp["type"])
            state = gt.minimize_blockmodel_dl(self.g, state_args=args)
            self.states.append({"state": state, "mdl": state.entropy()})

            mdl = self.states[-1]["mdl"]
            self.mdl_history.append(mdl)
            print(f"[{i+1}/{n_init}] MDL = {mdl:.2f}")

        best = min(self.states, key=lambda s: s["mdl"])
        self.mdl = best["mdl"]

    def fit_nested(self, n_init=10):
        """Fit nested SBM multiple times and keep best result."""
        self.states = []
        self.mdl_history = []

        for i in range(n_init):
            state = gt.minimize_nested_blockmodel_dl(self.g,
                        state_args=dict(deg_corr=self.deg_corr, clabel=self.g.vp["type"]))
            self.states.append({"state": state, "mdl": state.entropy()})
            self.mdl_history.append(state.entropy())
            print(f"[{i+1}/{n_init}] MDL = {state.entropy():.2f}")

        best = min(self.states, key=lambda s: s["mdl"])
        self.mdl = best["mdl"]

    def mdls(self):
        """Return list of all MDLs from runs."""
        return [s["mdl"] for s in self.states]

    def save_model(self, path):
        """Save current model to disk."""
        with open(path, "wb") as f:
            pickle.dump(self, f)
        print("Saved model to", path)

    def load_model(self, path):
        """Load a model from disk into self."""
        with open(path, "rb") as f:
            model = pickle.load(f)
        self.__dict__.update(model.__dict__)
        print("Loaded model from", path)

    def plot_mdl_history(self):
        """Plot MDL values from all fitting attempts."""
        if not self.mdl_history:
            print("No MDL history to plot.")
            return
        plt.plot(self.mdl_history, marker="o")
        plt.title("MDL per run")
        plt.xlabel("Run")
        plt.ylabel("Minimum Description Length")
        plt.tight_layout()
        plt.show()