FASTEN-cli 1.0.0__py2.py3-none-any.whl

FASTEN/model.py

@@ -0,0 +1,152 @@
+from .common import np, torch, nn, F, os, json
+from .config import ModelArgs, ModelInput, ModelOutput
+from .param import ModelParam, Constraint
+from .utils import Scaler, Encoder
+from pydantic import ValidationError
+from zipfile import ZipFile
+import joblib, random
+
+
+class Network(nn.Module):
+    def __init__(self, args: ModelArgs, input_size: int, output_size: int):
+        super(Network, self).__init__()
+        self.layers, sizes = nn.ModuleList(), [input_size]
+        if not args.hidden_layers: sizes.append(output_size)
+        elif args.architecture == "pyramidal":
+            step = (args.hidden_size - output_size) / args.hidden_layers
+            for i in range(args.hidden_layers):
+                current_size = int(args.hidden_size - (i * step))
+                sizes.append(max(current_size, output_size) )
+            sizes.append(output_size)
+        elif args.architecture == "rectangular":
+            sizes.extend([args.hidden_size] * args.hidden_layers)
+            sizes.append(output_size)
+        for i in range(len(sizes) - 1):
+            self.layers.append(nn.Linear(sizes[i], sizes[i+1]))
+
+    def load_constraints(self, params: dict[str, ModelParam]):
+        for rule in Constraint.RULES: 
+            self.register_buffer(rule, torch.zeros(len(params), dtype = bool))
+        for value in Constraint.VALUES: 
+            self.register_buffer(value, torch.zeros(len(params), dtype = float))
+        for i, param in enumerate(params.values()):
+            for rule in Constraint.RULES: self.get_buffer(rule)[i] = param.priors.get_rule(rule)
+            for value in Constraint.VALUES: self.get_buffer(value)[i] = param.priors.get_value(value)
+        for label in ["low", "high"]:
+            loop = (param.base == "label" for param in params.values())
+            mask = np.fromiter(loop, dtype = bool, count = len(params))
+            self.register_buffer(label, torch.from_numpy(mask))
+        self.to(next(self.parameters()).device)
+
+    def load_scaler(self, scaler: Scaler):
+        if not scaler.fitted: return
+        self.register_buffer("min", torch.from_numpy(scaler.min))
+        self.register_buffer("range", torch.from_numpy(scaler.range))
+        self.to(next(self.parameters()).device)
+    
+    def forward(self, x): 
+        for layer in self.layers[:-1]:
+            x = F.relu(layer(x))
+        x = self.layers[-1](x)
+        x[:, self.greater_than] = F.softplus(x[:, self.greater_than]) + (self.lower[self.greater_than] - self.min[self.greater_than]) / self.range[self.greater_than]
+        x[:, self.less_than] = (self.upper[self.less_than] - self.min[self.less_than]) / self.range[self.less_than] - F.softplus(x[:, self.less_than])
+        x[:, self.between] = F.sigmoid(x[:, self.between]) * self.interval[self.between] / self.range[self.between] + (self.lower[self.between] - self.min[self.between]) / self.range[self.between]
+        x[:, self.high] = x[:, self.low] + F.softplus(x[:, self.high] - x[:, self.low]) 
+        return x
+
+
+class Model():
+    def __init__(self, config_file: str = None, model_file: str = None):
+        torch.set_default_dtype(torch.float64)
+        self.inputs: dict[str, ModelInput] = dict()
+        self.outputs: dict[str, ModelOutput] = dict()
+        self.params: dict[str, ModelParam] = dict()
+        self.args: ModelArgs = None
+        if config_file:
+            with open(config_file, "r") as file: 
+                self.config = json.load(file)
+        else: self.config = {"train": {}}
+        if not model_file: self.create()
+        else: self.load(model_file)
+
+    def validate_data(self, inputs: dict, outputs: dict):
+        for label, config in inputs.items():
+            try: self.inputs[label] = ModelInput(**config, label = label)
+            except ValidationError as error: raise error
+        for label, config in outputs.items():
+            try: output = ModelOutput(**config, label = label)
+            except ValidationError as error: raise error
+            params = output.dist.load_parameters(output)
+            self.params.update(params)
+            self.outputs[label] = output
+
+    def validate_args(self, model: dict, train: dict):
+        model = {k: v for k, v in model.items() if not isinstance(v, list)}
+        train = {k: v for k, v in train.items() if not isinstance(v, list)}
+        try: self.args = ModelArgs(**{**model, **train})
+        except ValidationError as error: raise error
+        if self.args.rand_seed is not None: 
+            self.set_seed(self.args.rand_seed)
+        self.network = Network(self.args, len(self.inputs), len(self.params))
+        self.network = self.network.to(self.args.device)
+        self.network.load_constraints(self.params)
+        self.network.load_scaler(self.param_scaler)
+
+    @staticmethod
+    def set_seed(seed: int):
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        os.environ['PYTHONHASHSEED'] = str(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        random.seed(seed)
+        np.random.seed(seed)
+    
+    def create(self):
+        self.input_scaler, self.param_scaler = Scaler(), Scaler()
+        self.validate_data(self.config["inputs"], self.config["outputs"])
+        self.validate_args(self.config["model"], self.config["train"])
+        self.encoder = Encoder(self.inputs)
+        
+    def load(self, model_file):
+        with ZipFile(model_file, "r") as model:
+            with model.open("input_scaler.pkl", "r") as file:
+                self.input_scaler = joblib.load(file)
+            with model.open("param_scaler.pkl", "r") as file:
+                self.param_scaler = joblib.load(file)
+            with model.open("encoder.pkl", "r") as file:
+                self.encoder = joblib.load(file)  
+            with model.open("config.json", "r") as file:
+                model_config = json.load(file)
+                self.config["inputs"] = model_config["inputs"]
+                self.config["outputs"] = model_config["outputs"]
+                self.config["model"] = model_config["model"]
+            self.validate_data(self.config["inputs"], self.config["outputs"])
+            self.validate_args(self.config["model"], self.config["train"])
+            with model.open("model.pth", "r") as file:
+                self.network.load_state_dict(torch.load(file))
+                self.network.load_scaler(self.param_scaler)
+            
+    def dump(self, output):
+        model_file = f"{output}/model.zip"
+        with ZipFile(model_file, "w") as model:
+            with model.open("input_scaler.pkl", "w") as file:
+                joblib.dump(self.input_scaler, file)
+            with model.open("param_scaler.pkl", "w") as file:
+                joblib.dump(self.param_scaler, file)
+            with model.open("encoder.pkl", "w") as file:
+                joblib.dump(self.encoder, file)  
+            with model.open("model.pth", "w") as file:
+                torch.save(self.network.state_dict(), file) 
+            config = json.dumps(self.config, indent = 4)
+            model.writestr("config.json", config)
+
+    # def extract_config(self) -> dict:
+    #     train, model = self.args.model_dump_json(), dict()
+    #     train["device"], train["optimizer"] = train["device"].type, train["optimizer"].__name__
+    #     for key in ["architecture", "hidden_layers", "hidden_size"]: model[key] = train.pop(key)
+    #     inputs = {label: value.model_dump_json() for label, value in self.inputs.items()}
+    #     outputs = {label: value.model_dump_json() for label, value in self.outputs.items()}
+    #     return {"inputs": inputs, "outputs": outputs, "train": train, "model": model}
+        

FASTEN/param.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+from .common import torch, pd
+from torch.distributions import *
+from typing import TYPE_CHECKING
+from copy import deepcopy
+
+if TYPE_CHECKING: 
+    from .config import ModelOutput
+
+
+class ModelParam():
+    def __init__(self, output: ModelOutput, label: str): 
+        self.base: str = label.replace("0", "").replace("1", "")
+        self.label: str = f"{output.label}_{label}"
+        self.name: str = f"{output.name} {label.replace('_', ' ').title()}"
+        self.priors = Constraint(output, self.base)
+        self.constraints = deepcopy(self.priors)
+
+    def load_constraints(self, support: Support, min_val: float, max_val: float):
+        if support.upper == self.base: self.constraints.add_lower(max_val)
+        if support.lower == self.base: self.constraints.add_upper(min_val)
+        self.constraints.add_intervals()
+
+
+class ModelDist(): 
+    TYPES = {"Bernoulli", "Beta", "Binomial", "Cauchy", "Chi2", "ContinuousBernoulli", "Exponential", 
+        "FisherSnedecor", "Gamma", "Geometric", "Gumbel", "HalfCauchy", "HalfNormal", "InverseGamma", 
+        "Laplace", "LogNormal", "NegativeBinomial", "Normal", "Pareto", "Poisson", "Uniform"}  
+
+    def __init__(self, name: str):
+        self.name: str = name
+        self.params: dict[str, ModelParam] = dict()
+        self.support = Support(self.name)
+        if hasattr(ModelDist, name): self.base = getattr(ModelDist, name)
+        else: self.base = getattr(torch.distributions, name)
+    
+    def load_parameters(self, output: ModelOutput) -> dict[str, ModelParam]:
+        base = getattr(torch.distributions, self.name)
+        if self.name == "Uniform": labels = ["low", "high"]
+        else: labels = list(base.arg_constraints.keys())
+        for label in labels:
+            if label == "probs": continue
+            param = ModelParam(output, label)
+            self.params[f"{output.label}_{label}"] = param
+        return self.params
+
+    @staticmethod
+    def Bernoulli(logits: float | torch.Tensor) -> Bernoulli:
+        return Bernoulli(logits = logits, validate_args = False)
+    @staticmethod
+    def ContinuousBernoulli(logits: float | torch.Tensor) -> ContinuousBernoulli:
+        return ContinuousBernoulli(logits = logits, validate_args = False)
+    @staticmethod
+    def Geometric(logits: float | torch.Tensor) -> Geometric:
+        return Geometric(logits = logits, validate_args = False)
+    @staticmethod
+    def Binomial(total_count: float | torch.Tensor, logits: float | torch.Tensor) -> Binomial:
+        if isinstance(total_count, torch.Tensor):
+            rounded_total_count = total_count + (total_count.round() - total_count).detach()
+            return Binomial(rounded_total_count, logits = logits, validate_args = False)
+        else: return Binomial(round(total_count), logits = logits, validate_args = False)
+    @staticmethod
+    def NegativeBinomial(total_count: int | torch.Tensor, logits: float | torch.Tensor) -> NegativeBinomial:
+        return NegativeBinomial(total_count, logits = logits, validate_args = False)
+
+
+class Constraint():
+    POSITIVE = {"concentration", "total_count", "scale", "df", "rate", "alpha"}
+    RULES = ["greater_than", "less_than", "between"]
+    VALUES = ["lower", "upper", "interval"]
+
+    def __init__(self, output: ModelOutput, base: str, eps: float = 1e-16):
+        self.greater_than = self.less_than = self.between = False
+        self.lower = self.upper = self.interval = 0.0
+        if base in self.POSITIVE: self.greater_than, self.lower = True, eps
+        if hasattr(output, f"{base}_min"): 
+            self.add_lower(getattr(output, f"{base}_min"))
+        if hasattr(output, f"{base}_max"): 
+            self.add_upper(getattr(output, f"{base}_max"))
+        self.add_intervals()
+
+    def add_upper(self, upper): self.less_than, self.upper = True, upper
+    def add_lower(self, lower): self.greater_than, self.lower = True, lower
+    def add_intervals(self):
+        if self.greater_than and self.less_than:
+            self.greater_than = self.less_than = False
+            self.between, self.interval = True, self.upper - self.lower
+    
+    def get_rule(self, rule: str) -> bool: return getattr(self, rule)
+    def get_value(self, value: str) -> float: return getattr(self, value)
+
+
+class Support:
+    DISCRETE = {"Bernoulli", "Binomial", "Geometric", "NegativeBinomial", "Poisson"}
+    POSITIVE = {"FisherSnedecor", "InverseGamma", "LogNormal"}
+    NON_NEGATIVE = {"Binomial", "Poisson", "Geometric", "NegativeBinomial", "Chi2", "Exponential", "Gamma", "HalfCauchy", "HalfNormal"}
+    PROBABILITY = {"Bernoulli", "Beta", "ContinuousBernoulli"}
+    DEPENDENT = {"Uniform", "Binomial", "Pareto"}
+
+    def __init__(self, name, eps = 1e-16):
+        self.lower = self.upper = None
+        self.discrete = (name in self.DISCRETE)
+        if name in self.POSITIVE: self.lower = eps
+        if name in self.NON_NEGATIVE: self.lower = 0.0
+        if name in self.PROBABILITY: self.lower, self.upper = 0.0, 1.0
+        self.dependent = (name in self.DEPENDENT)
+        if name == "Uniform": self.lower, self.upper = "low", "high"
+        if name == "Binomial": self.upper = "total_count"
+        if name == "Pareto": self.lower = "scale"
+
+    def validate(self, data: pd.Series) -> bool:
+        if isinstance(self.lower, float) and (data < self.lower).any(): return False
+        if isinstance(self.upper, float) and (data > self.upper).any(): return False
+        return True
+
+    def get_bound(self, attr: str, params: dict[str, ModelParam], stats: torch.Tensor, fit: Distribution, std_devs: int = 3) -> torch.Tensor:
+        bound, sign = getattr(self, attr), -1 if attr == "lower" else 1
+        if bound is None: return (fit.mean + sign * std_devs * torch.sqrt(fit.variance)).unsqueeze(-1)
+        if isinstance(bound, str): return stats[:, [param.base == bound for param in params.values()]]
+        else: return torch.tensor([bound]).repeat(stats.shape[0]).unsqueeze(-1)

FASTEN/plot.py

@@ -0,0 +1,215 @@
+from .common import torch, pd, np, os
+import matplotlib.pyplot as plt
+import optuna
+
+COLORS = {"logits": "C6", "total_count": "C4", "rate": "C3",
+          "concentration": "C2", "loc": "C0", "scale": "C1",
+          "df": "C7", "alpha": "C5", "high": "C8", "low": "C9"}
+
+def plot_train(trainer, figure_dir):
+    if not os.path.exists(figure_dir): os.mkdir(figure_dir)
+    plt.figure(figsize = (6, 4))
+    plt.plot(range(len(trainer.train.loss)), trainer.train.loss, 
+        color = "blue", alpha = 0.5, label = "Training", zorder = 2)
+    max_val = max(trainer.train.loss)
+    if trainer.valid: 
+        plt.plot(range(len(trainer.valid.loss)), trainer.valid.loss, 
+            color = "red", alpha = 0.5, label = "Validation", zorder = 2)
+        max_val = max(max_val, max(trainer.valid.loss))
+    if max_val > 10: plt.yscale("log")
+    plt.xlabel("Epoch")
+    plt.legend()
+    plt.ylabel("Average Loss")
+    for i in range(2, 4):
+        plt.subplot(1, 3, i)
+        plt.plot([], [])
+        plt.axis("off")
+    plt.tight_layout()
+    plt.savefig(f"{figure_dir}/loss_curve_plot.png", dpi = 200)
+    plt.close()
+
+def plot_tune(tuner, figure_dir):
+    plt.figure(figsize = (4,4))
+    rank = optuna.importance.get_param_importances(tuner.study)
+    importances = np.array(list(rank.values()))
+    params = np.array(list(rank.keys()))
+    index = np.argsort(importances)
+    plt.barh(params[index], importances[index])
+    plt.xlabel("Importance")
+    plt.ylabel("Hyperparameter")
+    plt.tight_layout()
+    plt.savefig(f"{figure_dir}/importance_plot.png", dpi = 200)
+    plt.close()
+
+    trials = tuner.study.trials_dataframe()
+    plt.figure(figsize = (3 * len(params), 3.5))
+    for i, param in enumerate(params):
+        plt.subplot(1, len(params), i+1)
+        for value, group in trials.groupby(f"params_{param}"):
+            x, y = np.repeat(str(value), len(group)), group["value"]
+            plt.scatter(x, y, alpha = 0.5, linewidths = 0, color = "C0")
+            if y.max() - y.min() > 100: plt.yscale("log")
+        plt.xlabel(param)
+        plt.ylabel("Loss")
+        plt.xticks(rotation = 45, ha = "right")
+    plt.tight_layout()
+    plt.savefig(f"{figure_dir}/slices_plot.png", dpi = 200)
+    plt.close()
+
+    best = [trials.loc[:i+1, "value"].min() for i in range(len(trials))]
+    plt.figure(figsize = (8, 4))
+    plt.scatter(range(len(trials)), trials["value"], color = "C0")
+    if trials["value"].max() - trials["value"].min() > 100: plt.yscale("log")
+    plt.plot(range(len(trials)), best, color = "C3")
+    plt.xlabel("Trial")
+    plt.ylabel("Loss")
+    plt.tight_layout()
+    plt.savefig(f"{figure_dir}/convergence_plot.png", dpi = 200)
+    plt.close()
+
+
+def plot_predict(predictor, figure_dir): 
+    if predictor.true is None: return None, None, None
+    predictor.execute()
+    mse, kld, nll = predictor.evaluate()
+    if not os.path.exists(figure_dir): 
+        os.mkdir(figure_dir)
+    plot_statistics(predictor, figure_dir)
+    plot_kl_divergence(predictor, kld, figure_dir)
+    plot_mean_squares(predictor, mse, figure_dir)
+    plot_samples(predictor, figure_dir)
+    return mse, kld, nll
+
+def plot_samples(predictor, figure_dir, points = 10000, n_samples = 10):
+    sample_dir = f"{figure_dir}/samples"
+    if not os.path.exists(sample_dir): os.mkdir(sample_dir) 
+    groups = predictor.true.samples.outputs.groupby(predictor.true.samples.group)
+    n_samples = min(n_samples, len(groups) // 5)
+    clusters = predictor.true.stats.cluster_data(n_clusters = n_samples)
+    samples = pd.Series(np.arange(len(groups))).groupby(clusters).sample()
+    rows, cols = n_samples, len(predictor.model.outputs)
+    true = torch.empty((cols, points, rows)), torch.empty((cols, points, rows))
+    pred = torch.empty((cols, points, rows)), torch.empty((cols, points, rows))
+    (x_true, y_true), (x_pred, y_pred) = true, pred
+    
+    for j, output in enumerate(predictor.model.outputs.values()):
+        mask = [param in output.dist.params for param in predictor.model.params]
+        for label, (x, y) in {"true": true, "pred": pred}.items():
+            dataset = getattr(predictor, label)
+            stats = dataset.stats.outputs.iloc[samples,mask]
+            params = torch.tensor(stats.values)
+            fit = output.dist.base(*params.unbind(dim = 1))
+            lower = output.dist.support.get_bound("lower", output.dist.params, params, fit)
+            upper = output.dist.support.get_bound("upper", output.dist.params, params, fit)
+            dtype = int if output.dist.support.discrete else float
+            x[j] = torch.lerp(lower, upper, torch.linspace(0, 1, points)).to(dtype).t()
+            y[j] = torch.exp(fit.log_prob(x[j]))
+    
+    for i, sample in enumerate(samples):
+        plt.figure(figsize = (2.5 * cols + 0.5, 2))
+        group = groups.get_group(sample)
+        for j, output in enumerate(predictor.model.outputs.values()):
+            plt.subplot(1, cols, j + 1)
+            plt.hist(group.iloc[:,j], density = True, color = "darkgray")
+            real_true, real_pred = ~torch.isinf(y_true[j,:,i]), ~torch.isinf(y_pred[j,:,i])
+            mask_true, mask_pred = real_true.clone(), real_pred.clone()
+            mask_true[real_true] = y_true[j,real_true,i] >= 0.01 * max(y_true[j,real_true,i])
+            mask_pred[real_pred] = y_pred[j,real_pred,i] >= 0.01 * max(y_pred[j,real_pred,i])
+            plt.plot(x_true[j,mask_true,i], y_true[j,mask_true,i], color = "green", 
+                alpha = 1, linewidth = 2, label = f"Emprical (N = {len(group)})")
+            plt.plot(x_pred[j,mask_pred,i], y_pred[j,mask_pred,i], color = "orange", 
+                alpha = 1, linewidth = 2, linestyle = "dashed", label = "Predicted")
+            if j == cols - 1: plt.legend(bbox_to_anchor = (1.05, 1), loc = "upper left")
+            plt.xlabel(output.name)
+            plt.ylabel("Density")
+            plt.xticks(fontsize = 8)
+            plt.yticks(fontsize = 8)
+        plt.tight_layout()
+        plt.savefig(f"{sample_dir}/sample_plot_{sample}.png", dpi = 200)
+        plt.close()
+
+def plot_statistics(predictor, figure_dir):
+    cols, rows = len(predictor.model.outputs), 2
+    groups = predictor.true.samples.outputs.groupby(predictor.true.samples.group)
+    plt.figure(figsize = (3 * cols, 3 * rows))
+
+    for j, output in enumerate(predictor.model.outputs.values()):
+        mask = [param in output.dist.params for param in predictor.model.params]
+        stats = predictor.pred.stats.outputs.iloc[:,mask]
+        params = torch.tensor(stats.values)
+        fit = output.dist.base(*params.unbind(dim = 1))
+        pred_means, pred_vars = fit.mean, fit.variance
+        true_means = torch.zeros(groups.ngroups)
+        true_vars = torch.zeros(groups.ngroups)
+        for i, group in groups:
+            true_means[i] = group[output.label].mean()
+            true_vars[i] = group[output.label].var()
+        
+        plt.subplot(rows, cols, j + 1) 
+        plt.scatter(true_means, pred_means, color = "C0")
+        max_val = max(pred_means.max().max(), true_means.max().max())
+        min_val = min(pred_means.min().min(), true_means.min().min())
+        delta = (max_val - min_val) * 0.1
+        max_val, min_val = max_val + delta, min_val - delta
+        plt.axline((min_val, min_val), slope = 1, color = "darkgray", linestyle = "--")
+        plt.xlim(min_val, max_val)
+        plt.ylim(min_val, max_val)
+        plt.ylabel(f"Predicted Mean")
+        plt.xlabel(f"Empirical Mean")
+        plt.title(output.name, pad = 10, fontweight = "bold")
+
+        plt.subplot(rows, cols, cols + j + 1) 
+        plt.scatter(true_vars, pred_vars, color = "C1")
+        max_val = max(pred_vars.max().max(), true_vars.max().max())
+        min_val = min(pred_vars.min().min(), true_vars.min().min())
+        delta = (max_val - min_val) * 0.1
+        max_val, min_val = max_val + delta, min_val - delta
+        plt.axline((min_val, min_val), slope = 1, color = "darkgray", linestyle = "--")
+        plt.xlim(min_val, max_val)
+        plt.ylim(min_val, max_val)
+        plt.ylabel(f"Predicted Variance")
+        plt.xlabel(f"Empirical Variance")
+    plt.tight_layout(rect = [0, 0, 0.99, 1])
+    plt.savefig(f"{figure_dir}/statistics_plot.png", dpi = 200)
+    plt.close()
+    
+def plot_kl_divergence(predictor, kld, figure_dir):
+    names = [output.name for output in predictor.model.outputs.values()]
+    positions = range(1, kld.shape[1] + 1)
+    plt.figure(figsize = (6, 6))
+    bplot = plt.boxplot(kld, whis = [0, 100], patch_artist = True)
+    for patch in bplot["boxes"]: patch.set_facecolor("darkgray")
+    for patch in bplot["medians"]: patch.set_color("black")
+    plt.ylim(-0.5, kld.max().max() * 5)
+    plt.ylabel("KL Divergence")
+    plt.xticks(positions, names, rotation = 45, ha = "right")
+    plt.yscale("symlog")
+    plt.tight_layout()
+    plt.savefig(f"{figure_dir}/kl_divergence_plot.png", dpi = 200)
+    plt.close()
+
+def plot_mean_squares(predictor, mse, figure_dir):
+    names = [output.name for output in predictor.model.outputs.values()]
+    colors = [COLORS[param.base] for param in predictor.model.params.values()]
+    labels = [param.base for param in predictor.model.params.values()]
+    positions = [0.5*i for i in range(mse.shape[1])]
+    ticks = [0 for i in range(len(predictor.model.outputs))]
+    for i, output in enumerate(predictor.model.outputs.values()):
+        for j, param in enumerate(predictor.model.params):
+            if param in output.dist.params: 
+                positions[j] += 0.5*i + 1
+                ticks[i] += positions[j] / len(output.dist.params)
+    
+    plt.figure(figsize = (10, 6))
+    bplot = plt.boxplot(mse, whis = [0, 100], widths = 0.5, positions = positions, patch_artist = True, label = labels)
+    for patch, color in zip(bplot["boxes"], colors): patch.set_facecolor(color)
+    for patch in bplot["medians"]: patch.set_color("black")
+    plt.ylabel("Mean Squared Error")
+    plt.xticks(ticks, names, rotation = 45, ha = "right")
+    plt.yscale("log")
+    handles, labels = plt.gca().get_legend_handles_labels()
+    by_label = dict(zip(labels, handles))
+    plt.legend(by_label.values(), by_label.keys(), bbox_to_anchor = (1.02, 1), loc = "upper left")
+    plt.tight_layout()
+    plt.savefig(f"{figure_dir}/mean_squares_plot.png", dpi = 200)
+    plt.close()

FASTEN/predict.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+from .common import pd, np, torch
+from .data import Samples, Statistics, Dataset
+from .learn import Loss
+from .model import Model
+
+
+class Predictor():
+    def __init__(self, model: Model, dataset: Dataset = Dataset()):
+        self.model: Model = model
+        self.true: Dataset = dataset
+        self.pred: Dataset = Dataset()
+        self.error: pd.DataFrame = None
+        self.criterion: Loss = Loss(model)
+    
+    def load_inputs(self, inputs_file: str, num_runs: int): 
+        self.true.stats.load_data(inputs_file, self.model.inputs, None)
+        self.true.stats.scale_inputs(self.model.input_scaler)
+        num_runs = max(num_runs, 1)
+        index = np.arange(self.true.stats.inputs.shape[0]).repeat(num_runs)
+        self.true.samples.inputs = self.true.stats.inputs.loc[index]
+        self.true.samples.inputs = self.true.samples.inputs.reset_index(drop = True)
+        self.true.samples.group_data()
+        
+    def dump_statistics(self, outputs_file: str): 
+        self.pred.stats.dump_data(outputs_file)
+    
+    def dump_samples(self, outputs_file: str):
+        self.pred.samples.dump_data(outputs_file)
+
+    def execute(self):
+        self.pred.stats = self.predict_stats()
+        self.pred.stats.unscale_inputs(self.model.input_scaler)
+        self.pred.stats.unscale_outputs(self.model.param_scaler)
+        if isinstance(self.true.stats.outputs, pd.DataFrame): 
+            self.true.stats.unscale_outputs(self.model.param_scaler)
+        self.pred.samples = self.predict_samples()
+
+    def evaluate(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+        dependent = [output.dist.support.dependent for output in self.model.outputs.values()]
+        mse, kld, nll = self.criterion.evaluate(self.true, self.pred, any(dependent))
+        mse = pd.DataFrame(mse.detach().cpu().numpy(), columns = self.model.params)
+        kld = pd.DataFrame(kld.detach().cpu().numpy(), columns = self.model.outputs)
+        if any(dependent): return mse, kld, None
+        nll = pd.DataFrame(nll.detach().cpu().numpy(), columns = self.model.outputs)
+        return mse, kld, nll
+
+    def predict_stats(self) -> Statistics:
+        self.model.network.eval()
+        inputs = self.true.stats.inputs.values
+        x = torch.tensor(inputs).to(self.model.args.device)
+        with torch.no_grad(): y = self.model.network(x)
+        outputs = y.detach().cpu().numpy()
+        params = pd.DataFrame(outputs, columns = self.model.params)
+        return Statistics(self.true.stats.inputs, params)
+    
+    def predict_samples(self) -> Samples:
+        outputs = pd.DataFrame()
+        index = self.true.samples.group.values
+        for label, output in self.model.outputs.items():
+            mask = [param in output.dist.params for param in self.model.params]
+            stats = self.pred.stats.outputs.iloc[index, mask]
+            params = torch.tensor(stats.values)
+            fit = output.dist.base(*params.unbind(dim = 1))
+            outputs[label] = fit.sample().numpy()
+        return Samples(self.true.samples.inputs, outputs)

FASTEN/train.py

@@ -0,0 +1,87 @@
+from .common import torch, os, shutil
+from .learn import *
+from .data import Dataset
+from .model import Model
+from rich import progress
+
+
+class Trainer():    
+    def __init__(self, model: Model):
+        self.model: Model = model
+        self.dataset: Dataset = Dataset()
+        self.train: Partition = None
+        self.valid: Partition = None
+        self.test: Partition = None
+
+    def load_data(self, samples_file: str):
+        self.dataset.load_samples(samples_file, self.model)
+        self.dataset.load_stats(self.model, self.model.args.estimator)
+        self.split_data()
+
+    def split_data(self):
+        self.train = Partition(self.dataset, self.model.args.loss_func)
+        test = self.train.dataset.split(self.model.args.test_split, self.model.args.rand_seed) 
+        split = self.model.args.valid_split / (1 - self.model.args.test_split)
+        valid = self.train.dataset.split(split, self.model.args.rand_seed)
+        if test: self.test = Partition(test, self.model.args.loss_func)
+        if valid: self.valid = Partition(valid, self.model.args.loss_func)
+
+    def load_args(self):
+        torch.set_num_threads(os.cpu_count())
+        for partition in [self.train, self.valid, self.test]:
+            if partition: partition.load(self.model.args.batch_size, self.model.args.device)
+        args = {"lr": self.model.args.learn_rate, "weight_decay": self.model.args.weight_decay}
+        if self.model.args.optimizer.__name__ == "SGD": args["momentum"] = self.model.args.momentum
+        self.optimizer = self.model.args.optimizer(self.model.network.parameters(), **args)
+        patience, min_delta = self.model.args.patience, self.model.args.min_delta
+        if self.model.args.early_stop: self.early_stop = EarlyStop(patience, min_delta)
+        self.criterion = Loss(self.model).to(self.model.args.device)
+
+    def execute(self): 
+        self.load_args()
+        loadbar = progress.Progress(progress.TextColumn("{task.description}"), progress.BarColumn(), 
+                                    progress.MofNCompleteColumn(), progress.TimeRemainingColumn())
+        with loadbar as load:
+            task = load.add_task("", total = self.model.args.num_epochs)
+            for _ in range(self.model.args.num_epochs):
+                if self.train: self.train_model()
+                if self.valid: self.test_model(valid = True)
+                if self.valid and self.model.args.early_stop:
+                    loss = self.valid.loss[-1]
+                    if self.early_stop(loss): break
+                load.update(task, advance = 1)
+        if self.test: self.test_model(valid = False)
+
+    def train_model(self):
+        self.model.network.train()
+        train_loss = 0
+        for x, y in self.train.dataloader:
+            self.optimizer.zero_grad()
+            y_pred = self.model.network(x)
+            try: loss = self.criterion(y_pred, y)
+            except ValueError: raise ValueError("Training diverged: loss is NaN (possible exploding gradients)")
+            train_loss += loss.item()
+            loss.backward()
+            self.optimizer.step()
+        train_loss /= len(self.train.dataloader)
+        self.train.loss.append(train_loss)
+
+    def test_model(self, valid: bool):
+        if valid: partition = self.valid
+        else: partition = self.test
+        self.model.network.eval()
+        test_loss = 0
+        with torch.no_grad():
+            for x, y in partition.dataloader:
+                y_pred = self.model.network(x)
+                loss = self.criterion(y_pred, y)
+                test_loss += loss.item()
+        test_loss /= len(partition.dataloader)
+        partition.loss.append(test_loss)
+
+    def dump_model(self, output):
+        if os.path.exists(output): 
+            shutil.rmtree(output)
+        os.mkdir(output)
+        os.mkdir(f"{output}/plots")
+        self.model.dump(output)