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

FASTEN/__init__.py

@@ -0,0 +1,7 @@
+from .model import Model
+from .train import Trainer
+from .predict import Predictor
+from .tune import Tuner
+
+__version__ = "0.1.0"
+__all__ = ["Model", "Trainer", "Tuner", "Predictor"]

FASTEN/cli.py

@@ -0,0 +1,95 @@
+from .model import Model
+from .train import Trainer
+from .predict import Predictor
+from .tune import Tuner
+from .plot import plot_train, plot_predict, plot_tune
+from .common import pd
+from rich.console import Console
+import argparse, time
+
+def parse_args():
+    superparser = argparse.ArgumentParser(description = "A flexible software framework to approximate computationally \
+                                                         intensive simulations using neural-network-based emulators",
+                                          formatter_class = argparse.ArgumentDefaultsHelpFormatter)
+    subparsers = superparser.add_subparsers(dest = "command")
+    train = subparsers.add_parser("train", description = "Trains emulator on simulation data", 
+                                  formatter_class = argparse.ArgumentDefaultsHelpFormatter)
+    train.add_argument("-c", "--config", required = True, help = "JSON file defining configuration parameters")
+    train.add_argument("-i", "--input", required = True, help = "TSV file with simulation data")
+    train.add_argument("-o", "--output", default = "outputs", help = "Folder to output model and figures")
+    train.add_argument("-m", "--model", default = None, help = "ZIP file containing initial model")
+    tune = subparsers.add_parser("tune", description = "Tunes hyperparameters for emulator on simulation data", 
+                                 formatter_class = argparse.ArgumentDefaultsHelpFormatter)
+    tune.add_argument("-c", "--config", required = True, help = "JSON file defining configuration parameters")
+    tune.add_argument("-i", "--input", required = True, help = "TSV file with simulation data")
+    tune.add_argument("-o", "--output", default = "outputs", help = "Folder to output optimal configs and figures")
+    tune.add_argument("-n", "--trials", type = int, default = 100, help = "Total number of optimation trials")
+    tune.add_argument("--unique", action = "store_true", help = "Prevents re-training with duplicate hyperparameter sets")
+    predict = subparsers.add_parser("predict", description = "Predicts simulation data from inputs with emulator", 
+                                    formatter_class = argparse.ArgumentDefaultsHelpFormatter)
+    predict.add_argument("-m", "--model", required = True, help = "ZIP file containing model")
+    predict.add_argument("-i", "--input", required = True, help = "TSV file with simulation inputs")
+    predict.add_argument("-o", "--output", default = "outputs.tsv", help = "TSV file to output predicted simulation data")
+    predict.add_argument("-n", "--runs", default = 0, type = int, help = "Number of simulation runs per input")
+    return superparser.parse_args()
+
+def train(args, console):
+    model = Model(config_file = args.config, 
+                  model_file = args.model)
+    trainer = Trainer(model)
+    console.log("Estimating distribution parameters...")
+    trainer.load_data(args.input)
+    console.log("Training neural network...")
+    trainer.execute()
+    console.log("Writing outputs...")
+    trainer.dump_model(args.output)
+
+    plot_train(trainer, f"{args.output}/plots/training")
+    predictor = Predictor(model, trainer.train.dataset)
+    mse, kld, _ = plot_predict(predictor, f"{args.output}/plots/training")
+    console.print(f"Average Training MSE = {mse.mean().mean():.3g}\nAverage Training KL Divergence = {kld.mean().mean():.3g}")
+    if not trainer.test: return
+    predictor = Predictor(model, trainer.test.dataset)
+    mse, kld, _ = plot_predict(predictor, f"{args.output}/plots/testing")
+    if isinstance(mse, pd.DataFrame): 
+        console.print(f"Average Testing MSE = {mse.mean().mean():.3g}\nAverage Testing KL Divergence = {kld.mean().mean():.3g}")
+
+def predict(args, console): 
+    model = Model(model_file = args.model)
+    console.log("Predicting outputs...")
+    predictor = Predictor(model)
+    predictor.load_inputs(args.input, args.runs)
+    predictor.execute()
+    
+    if args.runs: predictor.dump_samples(args.output)
+    else: predictor.dump_statistics(args.output)
+
+def tune(args, console): 
+    model = Model(config_file = args.config)
+    trainer = Trainer(model)
+    if args.trials:
+        console.log("Estimating distribution parameters...")
+        trainer.load_data(args.input)
+
+    console.log("Tuning hyperparameters...")
+    tuner = Tuner(trainer)
+    tuner.load_study(args.output)
+    tuner.execute(args.trials, args.unique)
+
+    console.log("Writing outputs...")
+    tuner.dump_trials(args.output)
+    plot_tune(tuner, f"{args.output}/plots")
+
+def main():
+    args = parse_args()
+    start = time.perf_counter()
+    console = Console()
+    match args.command:
+        case "train": train(args, console)
+        case "predict": predict(args, console)
+        case "tune": tune(args, console)
+    end = time.perf_counter()
+    if end - start < 60: runtime = f"{(end - start):.2f} s"
+    elif end - start < 60 * 60: runtime = f"{(end - start) / 60:.2f} m"
+    else: runtime = f"{(end - start) / (60 * 60):.2f} h"
+    console.log(f"Done in {runtime}")

FASTEN/common.py

@@ -0,0 +1,5 @@
+import pandas as pd
+import numpy as np
+import torch, os, shutil, json
+import torch.nn as nn
+import torch.nn.functional as F

FASTEN/config.py

@@ -0,0 +1,110 @@
+from .common import pd, torch
+from .param import ModelDist
+from typing import Literal
+from warnings import warn
+import pydantic as pdc
+
+
+class ModelArgs(pdc.BaseModel):
+    test_split: float = pdc.Field(default = 0.1, ge = 0.0, lt = 1.0)
+    valid_split: float = pdc.Field(default = 0.1, ge = 0.0, lt = 1.0)
+    estimator: Literal["MoM", "MLE"] = "MLE"
+    rand_seed: int | None = None 
+
+    architecture: Literal["rectangular", "pyramidal"] = "pyramidal"
+    hidden_layers: int = pdc.Field(default = 2, ge = 0)
+    hidden_size: int = pdc.Field(default = 64, gt = 0)
+
+    device: Literal["cpu", "cuda"] = "cpu"    
+    batch_size: int = pdc.Field(default = 32, gt = 0)
+    num_epochs: int = pdc.Field(default = 1e5, gt = 0)
+
+    early_stop: bool = True
+    patience: int = pdc.Field(default = 20, ge = 0)
+    min_delta: float = pdc.Field(default = 0.0)
+    
+    optimizer: Literal["SGD", "Adam", "AdamW"] = "AdamW"
+    loss_func: Literal["MSE", "KLD", "NLL"] = "NLL"
+    learn_rate: float = pdc.Field(default = 1e-3, gt = 0.0) 
+    weight_decay: float = pdc.Field(default = 0.0, ge = 0.0)
+    momentum: float = pdc.Field(default = 0.0, ge = 0.0)
+
+    @pdc.field_validator("device", mode = "after")
+    @classmethod
+    def validate_device(cls, value):
+        if value == "cuda" and not torch.cuda.is_available():
+            warn("PyTorch cannot find a compatible GPU. Defaulting to CPU.")
+            return torch.device("cpu")
+        return torch.device(value)
+    
+    @pdc.field_validator("optimizer", mode = "after")
+    @classmethod
+    def validate_optimizer(cls, value: str):
+        return getattr(torch.optim, value)
+    
+    @pdc.model_validator(mode = "after")
+    def validate_early_stop(self):
+        if not self.valid_split and self.early_stop:
+            raise ValueError("Non-empty validation set required for early stopping.")
+        return self
+    
+    @pdc.model_validator(mode = "after")
+    def validate_splits(self):
+        if self.valid_split + self.test_split >= 1:
+            raise ValueError("Non-empty training set required. Decrease size of validation or testing set.")
+        return self
+
+
+class ModelInput(pdc.BaseModel):
+    label: str
+    name: str = pdc.Field(default_factory = lambda data: data['label'])
+    type: Literal["float", "integer", "string"] = "float"
+
+    def validate_data(self, data: pd.DataFrame, label: str):
+        if data[label].isna().any():
+            raise ValueError(f"Training data contains missing or undefined values: {self.name}.")
+        if self.type == "string" and not pd.api.types.is_string_dtype(data[label]):
+            raise ValueError(f"Training data has invalid values: {self.name}.")
+        if self.type in ["integer", "float"]:
+            if not pd.api.types.is_numeric_dtype(data[label]):
+                raise ValueError(f"Training data has invalid values: {self.name}.")
+            else: data[label] = data[label].astype(float)
+        if self.type == "integer" and (data[label] % 1 != 0).any(): 
+            warn(f"Integer type specified for non-integer training data: {self.name}. Rounding to nearest integer.")
+            data[label] = data[label].round()
+
+
+class ModelOutput(pdc.BaseModel): # validate priors
+    model_config = pdc.ConfigDict(arbitrary_types_allowed = True, extra = "allow")
+
+    label: str
+    dist: str | ModelDist
+    name: str = pdc.Field(default_factory = lambda data: data['label'])
+    type: Literal["float", "integer"] = "float"
+    min_thresh: float | None = None
+    max_thresh: float | None = None
+
+    @pdc.field_validator("dist", mode = "after")
+    @classmethod
+    def validate_distribution(cls, value: str) -> ModelDist:
+        try: dist = ModelDist(value)
+        except AttributeError: raise ValueError(f"Invalid distribution specified.")
+        return dist
+    
+    @pdc.model_validator(mode = "after")
+    def validate_discrete(self):
+        if self.dist.support.discrete and self.type != "integer":
+            raise ValueError(f"Discrete distribution specified for non-integer training data: {self.name}.")
+        return self
+    
+    def validate_data(self, data: pd.DataFrame, label: str):
+        if data[label].isna().any():
+            raise ValueError(f"Training data contains missing or undefined values: {self.name}.")
+        if not pd.api.types.is_numeric_dtype(data[label]):
+            raise ValueError(f"Training data has invalid values: {self.name}.")
+        else: data[label] = data[label].astype(float)
+        if self.type == "integer" and (data[label] % 1 != 0).any(): 
+            warn(f"Integer type specified for non-integer training data: {self.name}. Rounding to nearest integer.")
+            data[label] = data[label].round()
+        if not self.dist.support.validate(data[label]): 
+            raise AssertionError(f"Training data contains values outside domain of distribution: {self.name}")

FASTEN/data.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+from .common import pd, os, np
+from .config import ModelOutput
+from .utils import Scaler, Encoder
+from .estimate import Estimator
+from .model import Model
+from sklearn.model_selection import train_test_split
+from k_means_constrained import KMeansConstrained
+from typing import Any
+from abc import ABC
+
+
+class Dataset():
+    def __init__(self, samples: Samples = None, stats: Statistics = None):
+        self.samples = samples if samples else Samples()
+        self.stats = stats if stats else Statistics()
+
+    def load_samples(self, data_file: str, model: Model):
+        self.samples.load_data(data_file, model.inputs, model.outputs)
+        self.samples.filter_data(model.outputs)
+        for label, config in model.outputs.items(): 
+            config.validate_data(self.samples.outputs, label)
+        for label, config in model.inputs.items(): 
+            config.validate_data(self.samples.inputs, label)
+        self.samples.encode_inputs(model.encoder)
+        self.samples.scale_inputs(model.input_scaler)
+        self.samples.group_data()
+
+    def load_stats(self, model: Model, estimator: str):
+        self.estimate_stats(Estimator(estimator, model.outputs))
+        self.stats.scale_outputs(model.param_scaler)
+        model.network.load_scaler(model.param_scaler)
+
+    def estimate_stats(self, estimator: Estimator):
+        input_data = self.samples.inputs.groupby(self.samples.group)
+        output_data = self.samples.outputs.groupby(self.samples.group)
+        self.stats.inputs = input_data.first().reset_index(drop = True)
+        self.stats.outputs = estimator.execute(output_data)
+
+    def split(self, split_prop: float, rand_seed: int): 
+        if not split_prop: return None
+        groups, index = self.stats.cluster_data(split_prop), self.stats.inputs.index
+        train_index, test_index = train_test_split(index, test_size = split_prop, stratify = groups, random_state = rand_seed)
+        train_index, test_index = sorted(train_index), sorted(test_index)
+        train_samples, test_samples = self.samples.group.isin(train_index), self.samples.group.isin(test_index)
+        test_stats = Statistics(self.stats.inputs.loc[test_index].reset_index(drop = True), 
+            None if self.stats.outputs is None else self.stats.outputs.loc[test_index].reset_index(drop = True))
+        self.stats = Statistics(self.stats.inputs.loc[train_index].reset_index(drop = True), 
+            None if self.stats.outputs is None else self.stats.outputs.loc[train_index].reset_index(drop = True))
+        test_samples = Samples(self.samples.inputs.loc[test_samples].reset_index(drop = True), 
+            self.samples.outputs.loc[test_samples].reset_index(drop = True))
+        self.samples = Samples(self.samples.inputs.loc[train_samples].reset_index(drop = True), 
+            self.samples.outputs.loc[train_samples].reset_index(drop = True))
+        return Dataset(samples = test_samples, stats = test_stats)
+
+
+class Data(ABC): 
+    def __init__(self, inputs: pd.DataFrame = None, outputs: pd.DataFrame = None):
+        self.inputs, self.outputs = inputs, outputs
+    
+    def _repr_html_(self): 
+        return pd.concat([self.inputs, self.outputs], axis = 1)._repr_html_()
+
+    def __str__(self): 
+        return pd.concat([self.inputs, self.outputs], axis = 1).__str__()
+
+    def dump_data(self, data_file: str):
+        data = pd.concat([self.inputs, self.outputs], axis = 1)
+        data.to_csv(data_file, sep = '\t', index = False)
+    
+    def load_data(self, data_file: str, inputs: dict[str, Any], outputs: dict[str, Any]):
+        inputs, outputs = list(inputs), list(outputs) if outputs else None
+        self.inputs = pd.read_csv(data_file, sep = "\t", usecols = inputs)[inputs]
+        self.inputs = self.inputs.sort_values(by = inputs)
+        if outputs: 
+            self.outputs = pd.read_csv(data_file, sep = "\t", usecols = outputs)[outputs]
+            self.outputs = self.outputs.loc[self.inputs.index].reset_index(drop = True)
+        self.inputs = self.inputs.reset_index(drop = True)
+
+
+class Samples(Data):
+    def __init__(self, inputs: pd.DataFrame = None, outputs: pd.DataFrame = None):
+        super().__init__(inputs, outputs)
+        if isinstance(inputs, pd.DataFrame): self.group_data()
+        else: self.group: pd.Series = None
+
+    def filter_data(self, outputs: dict[str, ModelOutput]):
+        for label, output in outputs.items():
+            if output.min_thresh is not None:
+                mask = (self.outputs[label] > output.min_thresh)
+                self.outputs = self.outputs[mask].reset_index(drop = True)
+                self.inputs = self.inputs[mask].reset_index(drop = True)
+            if output.max_thresh is not None:
+                mask = (self.outputs[label] < output.max_thresh)
+                self.outputs = self.outputs[mask].reset_index(drop = True)
+                self.inputs = self.inputs[mask].reset_index(drop = True)
+
+    def group_data(self):
+        matches = (self.inputs != self.inputs.shift())
+        self.group = matches.any(axis = 1).cumsum() - 1
+
+    def scale_inputs(self, scaler: Scaler): scaler.transform(self.inputs)
+    def unscale_inputs(self, scaler: Scaler): scaler.inverse_transform(self.inputs)
+    def encode_inputs(self, encoder: Encoder): encoder.transform(self.inputs)
+
+
+class Statistics(Data): 
+    def cluster_data(self, split_prop: float = 1, n_clusters: int | None = None) -> np.ndarray:
+        if n_clusters is None:
+            n_clusters = int(self.inputs.shape[0] * split_prop / 5)
+        if n_clusters < 1: return None
+        kmeans = KMeansConstrained(n_clusters = n_clusters, size_min = 5)
+        return kmeans.fit_predict(self.inputs)
+    
+    def scale_outputs(self, scaler: Scaler): scaler.transform(self.outputs)
+    def unscale_outputs(self, scaler: Scaler): scaler.inverse_transform(self.outputs)
+    def scale_inputs(self, scaler: Scaler): scaler.transform(self.inputs)
+    def unscale_inputs(self, scaler: Scaler): scaler.inverse_transform(self.inputs)

FASTEN/estimate.py

@@ -0,0 +1,138 @@
+from .common import np, pd, torch, F, os
+from .param import ModelDist, Constraint
+from .config import ModelOutput
+from pandas.api.typing import DataFrameGroupBy
+from rich import progress
+
+PROGRESS = progress.Progress(
+    progress.TextColumn("{task.description}"),
+    progress.BarColumn(),
+    progress.MofNCompleteColumn(),
+    progress.TimeRemainingColumn(),
+)
+
+class Estimator(): 
+
+    class Moments():
+        @staticmethod
+        def Exponential(data: pd.Series) -> torch.Tensor:
+            mean = data.mean()
+            if not mean: raise AssertionError()
+            return torch.tensor([1 / mean])
+        @staticmethod
+        def Normal(data: pd.Series) -> torch.Tensor: 
+            mean, var = data.mean(), data.var()
+            if np.isnan(var) or not var: raise AssertionError()
+            return torch.tensor([mean, np.sqrt(var)])
+        @staticmethod
+        def HalfNormal(data: pd.Series) -> torch.Tensor: 
+            mean = np.sqrt(data.pow(2).mean())
+            if not mean: raise AssertionError()
+            return torch.tensor([mean])
+        @staticmethod
+        def LogNormal(data: pd.Series) -> torch.Tensor: 
+            mean, var = np.log(data).mean(), np.sqrt(np.log(data).var())
+            if np.isnan(var) or not var: raise AssertionError()
+            return torch.tensor([mean, var])
+        @staticmethod
+        def Uniform(data: pd.Series) -> torch.Tensor: 
+            if data.min() == data.max(): # raise AssertionError()
+                return torch.tensor([data.min(), data.max()])
+            return torch.tensor([data.min(), data.max()])
+        @staticmethod
+        def Geometric(data: pd.Series) -> torch.Tensor: 
+            return torch.tensor([1 / (1 + data.mean())])
+        @staticmethod
+        def Poisson(data: pd.Series) -> torch.Tensor: 
+            mean = data.mean()
+            if not mean: raise AssertionError()
+            return torch.tensor([mean])
+        @staticmethod
+        def Bernoulli(data: pd.Series) -> torch.Tensor: 
+            return torch.tensor([data.mean()])
+        @staticmethod
+        def Laplace(data: pd.Series) -> torch.Tensor:
+            mad = (data - data.median()).abs().mean() 
+            if mad <= 0: raise AssertionError()
+            return torch.tensor([data.median(), mad])
+        @staticmethod
+        def Pareto(data: pd.Series) -> torch.Tensor: 
+            log_sum = np.log(data / data.min()).sum()
+            if not log_sum: raise AssertionError()
+            return torch.tensor([data.min(), data.shape[0] / log_sum])
+        @staticmethod
+        def Binomial(data: pd.Series) -> torch.Tensor: 
+            mean, var = data.mean(), data.var()
+            if not var or np.isnan(var): raise AssertionError()
+            if not mean or var >= mean: raise AssertionError()
+            probs = 1 - var / mean
+            total_counts = max(mean / probs, data.max())
+            return torch.tensor([total_counts, np.log(probs / (1 - probs))])
+        @staticmethod
+        def NegativeBinomial(data: pd.Series) -> torch.Tensor:
+            mean, var = data.mean().item(), data.var().item()
+            if not var or np.isnan(var): raise AssertionError()
+            if not mean or var <= mean: raise AssertionError()
+            total_counts, probs = mean**2 / (var - mean), 1 - mean / var
+            return torch.tensor([total_counts, np.log(probs / (1 - probs))])
+    
+    def __init__(self, estimator: str, outputs: dict[str, ModelOutput]):
+        self.estimator: str = estimator
+        self.outputs: dict[str, ModelOutput] = outputs
+
+    def execute(self, groups: DataFrameGroupBy) -> pd.DataFrame:
+        torch.set_num_threads(os.cpu_count())
+        with PROGRESS as progress:
+            params = [self.iterate(groups, label, progress) for label in self.outputs]
+        return pd.concat(params, axis = 1)
+    
+    def iterate(self, total_groups: DataFrameGroupBy, label: str, progress: progress.Progress) -> pd.DataFrame:
+        output, groups, params = self.outputs[label], total_groups[label], dict()
+        task = progress.add_task(output.name, total = groups.ngroups)
+        for group, data in groups:
+            params[group] = self.estimate(output, data)
+            progress.update(task, advance = 1)
+        return pd.DataFrame.from_dict(params, "index", None, output.dist.params).sort_index()
+        
+    def estimate(self, output: ModelOutput, data: pd.Series) -> np.ndarray: 
+        if self.estimator == "MoM" and hasattr(self.Moments, output.dist.name):
+            method = getattr(self.Moments, output.dist.name)
+            try: return method(data).numpy()
+            except AssertionError: pass
+        return self.max_likelihood(output.dist, torch.tensor(data.values))
+
+    def max_likelihood(self, dist: ModelDist, data: torch.Tensor) -> np.ndarray: 
+        self.load_constraints(dist, data)
+        if dist.support.discrete: data = data.to(int)
+        weights = torch.randn(len(dist.params), requires_grad = True, dtype = float)
+        optimizer = torch.optim.LBFGS([weights], max_iter = 200, line_search_fn = "strong_wolfe",
+                                      tolerance_grad = 1e-12, tolerance_change = 1e-12)
+
+        def closure():
+            optimizer.zero_grad()
+            params = self.apply_constraints(weights) 
+            params.nan_to_num_(nan = 1e-16)
+            fit = dist.base(*params)
+            loss = -1 * fit.log_prob(data).mean()
+            loss.backward()
+            return loss
+
+        optimizer.step(closure)
+        with torch.no_grad(): 
+            return self.apply_constraints(weights).numpy()
+
+    def load_constraints(self, dist: ModelDist, data: pd.Series): 
+        for rule in Constraint.RULES: setattr(self, rule, torch.zeros(len(dist.params), dtype = bool))
+        for value in Constraint.VALUES: setattr(self, value, torch.zeros(len(dist.params), dtype = float))
+        for i, param in enumerate(dist.params.values()):
+            min_val, max_val = data.min().item(), data.max().item()
+            param.load_constraints(dist.support, min_val, max_val)
+            for rule in Constraint.RULES: getattr(self, rule)[i] = param.constraints.get_rule(rule)
+            for value in Constraint.VALUES: getattr(self, value)[i] = param.constraints.get_value(value)
+
+    def apply_constraints(self, weights: torch.Tensor) -> torch.Tensor:
+        params = weights.clone()
+        params[self.greater_than] = F.softplus(params[self.greater_than]) + self.lower[self.greater_than]
+        params[self.less_than] = self.upper[self.less_than] - F.softplus(params[self.less_than])
+        params[self.between] = F.sigmoid(params[self.between]) * self.interval[self.between] + self.lower[self.between]
+        return params

FASTEN/learn.py

@@ -0,0 +1,165 @@
+from .common import torch, nn, F, np
+from .data import Dataset
+from .model import Model
+from torch.distributions import NegativeBinomial, Binomial
+from torch.distributions.kl import register_kl
+
+
+MAX_STEPS = 100000
+MAX_SIZE = 4 * 1024**2
+
+@register_kl(NegativeBinomial, NegativeBinomial)
+def KL_negative_binomial(p, q): 
+    log_p, log_q = F.logsigmoid(p.logits), F.logsigmoid(q.logits)
+    log_neg_p, log_neg_q = F.logsigmoid(-p.logits), F.logsigmoid(-q.logits)
+    d_log, d_log_neg = log_p - log_q, log_neg_p - log_neg_q
+    mean_p = p.total_count * torch.exp(p.logits)
+    kld_exact = mean_p * d_log + p.total_count * d_log_neg
+    log_mean_p = torch.log(p.total_count) + p.logits
+    std_p = torch.exp((log_mean_p - log_neg_p) / 2)
+    return approximate_KL(p, q, mean_p, std_p, kld_exact)
+
+@register_kl(Binomial, Binomial)
+def KL_binomial(p, q): 
+    log_p, log_q = F.logsigmoid(p.logits), F.logsigmoid(q.logits)
+    log_neg_p, log_neg_q = F.logsigmoid(-p.logits), F.logsigmoid(-q.logits)
+    d_log, d_log_neg = log_p - log_q, log_neg_p - log_neg_q
+    mean_p = p.total_count * torch.sigmoid(p.logits)
+    kld_exact = mean_p * d_log + (p.total_count - mean_p) * d_log_neg
+    std_p = torch.sqrt(mean_p * (1 - torch.sigmoid(p.logits)))
+    return approximate_KL(p, q, mean_p, std_p, kld_exact)
+
+def approximate_KL(p, q, mean_p, std_p, kld_exact, n_stds = 6):
+    comparable = torch.isclose(p.total_count, q.total_count)
+    if comparable.all(): return kld_exact
+
+    max_k = torch.ceil(mean_p + n_stds * std_p)
+    min_k = torch.clamp(torch.floor(mean_p - n_stds * std_p), min = 0.0)
+    zeros, ones = torch.zeros_like(std_p), torch.ones_like(std_p)
+    total_size = torch.where(comparable, zeros, max_k - min_k)
+    integer_size = torch.floor(total_size) + 1
+    n_steps = int(max(1, min(integer_size.max().item(), MAX_STEPS)))
+
+    max_width = total_size / max(n_steps - 1, 1)
+    width = torch.where(total_size > n_steps, max_width, ones)
+    full = torch.full_like(total_size, n_steps)
+    valid = torch.where(total_size > n_steps, full, integer_size)
+    size = int(max(1, min(MAX_SIZE, n_steps)))
+
+    kld_approx = torch.zeros_like(mean_p)
+    min_k, valid = min_k.unsqueeze(0), valid.unsqueeze(0)
+    width = width.unsqueeze(0)
+    for i in range(0, n_steps, size):
+        j = min(i + size, n_steps)
+        shape = [-1] + [1] * len(mean_p.shape)
+        delta = torch.arange(i, j, device = p.logits.device).view(shape)
+        k = torch.where(delta < valid, min_k + delta * width, min_k)
+        log_prob_p, log_prob_q = p.log_prob(k), q.log_prob(k)        
+        kld = torch.exp(log_prob_p) * (log_prob_p - log_prob_q) * width
+        kld = torch.where(delta < valid, kld, torch.zeros_like(kld))
+        kld_approx += torch.sum(kld, dim = 0)
+    return torch.where(comparable, kld_exact, kld_approx)
+
+
+class Partition(): 
+    def __init__(self, dataset: Dataset, loss_func: str):
+        self.dataset: Dataset = dataset
+        self.dataloader: torch.utils.data.DataLoader = None
+        self.by_sample: bool = (loss_func == "NLL")
+        self.loss: list = []
+
+    def load(self, batch_size: int, device: torch.device): 
+        data = self.dataset.samples if self.by_sample else self.dataset.stats
+        input_tensor = torch.tensor(data.inputs.values).to(device)
+        output_tensor = torch.tensor(data.outputs.values).to(device)
+        data_tensor = torch.utils.data.TensorDataset(input_tensor, output_tensor)
+        self.dataloader = torch.utils.data.DataLoader(data_tensor, batch_size, shuffle = True)
+        self.loss = []
+
+
+class EarlyStop:
+    def __init__(self, patience = 20, min_delta = 0, multiplier = 0.1):
+        self.multiplier: float = multiplier
+        self.patience: float = patience
+        self.min_delta: float = min_delta
+        self.counter: int = 0
+        self.best_loss: float = float('inf')
+        self.avg_loss: float = None
+
+    def __call__(self, loss):
+        if self.avg_loss is not None: 
+            self.avg_loss *= 1 - self.multiplier
+            self.avg_loss += self.multiplier * loss
+        else: self.avg_loss = loss
+        if self.avg_loss < self.best_loss - self.min_delta:
+            if self.avg_loss < self.best_loss:
+                self.best_loss = self.avg_loss
+            self.counter = 0
+        else: self.counter += 1
+        return (self.counter >= self.patience)
+    
+
+class Loss(nn.Module):
+    def __init__(self, model: Model):
+        super().__init__()
+        self.loss_func = model.args.loss_func
+        self.mean_squared_error = nn.MSELoss(reduction = "none")
+        self.register_buffer("min", torch.from_numpy(model.param_scaler.min))
+        self.register_buffer("range", torch.from_numpy(model.param_scaler.range))
+        self.load_params(model)
+
+    def load_params(self, model: Model):
+        self.num_outputs = len(model.outputs)
+        self.dists = [output.dist.base for output in model.outputs.values()]
+        masks = torch.zeros((self.num_outputs, len(model.params)), dtype = bool)
+        for i, output in enumerate(model.outputs.values()): 
+            loop = (param in output.dist.params for param in model.params)
+            mask = np.fromiter(loop, dtype = bool, count = len(model.params))
+            masks[i] = torch.tensor(mask, dtype = bool)
+        self.register_buffer("outputs", masks)
+
+    def forward(self, pred: torch.Tensor, true: torch.Tensor):
+        pred_scaled, true_scaled = pred, true
+        pred_unscaled = pred * self.range + self.min
+        if self.loss_func != "NLL":
+            true_unscaled = true * self.range + self.min
+        match self.loss_func:
+            case "MSE": return self.mean_squared_error(pred_scaled, true_scaled).mean()
+            case "KLD": return self.kl_divergence(pred_unscaled, true_unscaled).mean()
+            case "NLL": return self.neg_log_likelihood(pred_unscaled, true_scaled).mean()
+
+    def evaluate(self, pred: Dataset, true: Dataset, dependent: bool) -> list[torch.Tensor, torch.Tensor, torch.Tensor]:
+        pred_stats_unscaled = torch.tensor(pred.stats.outputs.values)
+        true_stats_unscaled = torch.tensor(true.stats.outputs.values)
+        pred_stats_scaled = (pred_stats_unscaled - self.min) / self.range
+        true_stats_scaled = (true_stats_unscaled - self.min) / self.range
+        mse = self.mean_squared_error(pred_stats_scaled, true_stats_scaled)
+        kld = self.kl_divergence(pred_stats_unscaled, true_stats_unscaled)
+        if dependent: return mse, kld, None
+        sample_groups = true.samples.outputs.groupby(true.samples.group)
+        nll = torch.zeros((sample_groups.ngroups, pred.samples.outputs.shape[1]), dtype = float)
+        for i, sample_group in sample_groups:
+            pred_stats = pred_stats_unscaled[i].unsqueeze(0).repeat(len(sample_group), 1)
+            true_samples = torch.tensor(sample_group.values)
+            nll[i] = self.neg_log_likelihood(pred_stats, true_samples)
+        return mse, kld, nll
+    
+    def neg_log_likelihood(self, pred: torch.Tensor, true: torch.Tensor) -> torch.Tensor:
+        neg_log_likelihoods = []
+        for i in range(self.num_outputs):
+            pred_params = pred[:, self.outputs[i]]
+            pred_fit = self.dists[i](*pred_params.unbind(dim = 1))
+            loss = -1 * pred_fit.log_prob(true[:,i])
+            neg_log_likelihoods.append(loss.mean())
+        return torch.column_stack(neg_log_likelihoods)
+
+    def kl_divergence(self, pred: torch.Tensor, true: torch.Tensor) -> torch.Tensor:
+        kl_divergences = []
+        for i in range(self.num_outputs):
+            pred_params = pred[:, self.outputs[i]]
+            true_params = true[:, self.outputs[i]]
+            pred_fit = self.dists[i](*pred_params.unbind(dim = 1))
+            true_fit = self.dists[i](*true_params.unbind(dim = 1))
+            loss = torch.distributions.kl_divergence(true_fit, pred_fit)
+            kl_divergences.append(loss)
+        return torch.column_stack(kl_divergences)