Repository of Conditioning Score-Based Generative Models by Neuro-Symbolic Constraints paper
from src.constants import MODELS_FOLDER
from src.constrained_generation_config import ConstrainedGenerationConfig
from src.constraints.constraint import Constraint
from src.constraints.gradient_mixer import *
from src.constraints.real_logic import Product as Logic
from src.models.score_based.sdes.sampler import EulerMethod
from src.util import find

constraint = Constraint(
    f=lambda x: Logic.or_(Logic.greater(x, 0.5), Logic.smaller(x, -0.5)),
    strength=1.0,
    gradient_mixer=linear_interpolation,
)

generation_options = dict(
    sde_solver=EulerMethod(),
    n_samples=4000,
    steps=1500,
    corrector_steps=0,
    final_corrector_steps=0,
)

CONFIG = ConstrainedGenerationConfig(
    name="not test",
    constraint=constraint,
    model_path=find(str(MODELS_FOLDER), pattern="*normal*"),  # type: ignore
    generation_options=generation_options,
)

config/constrained_generation/debug.py
DavideScassola/score-based-constrained-generation
from torch import Tensor

from src.constants import MODELS_FOLDER
from src.constrained_generation_config import ConstrainedGenerationConfig
from src.constraints.constraint import Constraint
from src.constraints.gradient_mixer import *
from src.constraints.real_logic import Product as Logic
from src.models.score_based.sdes.sampler import EulerMethod
from src.util import find, names_index_map

NAME = "eSIRS"

S = 0
I = 1

GRAD = 7.0
SEED = 1234
SCORE_APPROXIMATOR = snr

positive_populations = lambda x: Logic.and_(Logic.greater(x, 0.0, grad=1))
limited_population = lambda x: Logic.and_(Logic.smaller(x.sum(dim=-1), 100.0, grad=1))
fix_S_at_0 = lambda x: Logic.equal(x[:, 0, S], 80.0, grad=GRAD)
fix_I_at_0 = lambda x: Logic.equal(x[:, 0, I], 10.0, grad=GRAD)
fix_S_at_25 = lambda x: Logic.equal(x[:, 25, S], 30.0, grad=GRAD)


def predicate(x: Tensor) -> Tensor:
    return Logic.and_(
        positive_populations(x),
        limited_population(x),
        fix_S_at_0(x),
        fix_I_at_0(x),
        fix_S_at_25(x),
    )


constraint = Constraint(
    f=lambda x: predicate(x), strength=1.0, gradient_mixer=SCORE_APPROXIMATOR
)

generation_options = dict(
    sde_solver=EulerMethod(),
    n_samples=5000,
    steps=1000,
    corrector_steps=2,
    final_corrector_steps=2000,
)

CONFIG = ConstrainedGenerationConfig(
    name=NAME + "_bridging",
    constraint=constraint,
    model_path=find(str(MODELS_FOLDER), pattern=f"*{NAME}*"),
    generation_options=generation_options,
    seed=SEED,
)

config/constrained_generation/eSIRS_bridging.py
DavideScassola/score-based-constrained-generation
from torch import Tensor

from src.constants import MODELS_FOLDER
from src.constrained_generation_config import ConstrainedGenerationConfig
from src.constraints.constraint import Constraint
from src.constraints.gradient_mixer import (linear_interpolation,
                                            noise_weighting, snr)
from src.constraints.real_logic import Product as Logic
from src.models.score_based.sdes.sampler import EulerMethod
from src.util import find

NAME = "eSIRS"

S = 0
I = 1

GRAD = 25.0
SEED = 1234

positive_populations = lambda x: Logic.and_(Logic.greater(x, 0.0, grad=1))
limited_population = lambda x: Logic.and_(Logic.smaller(x.sum(dim=2), 100.0, grad=1))
limited_I_from_0 = lambda x: Logic.and_(Logic.smaller(x[:, 0:, I], 20.0, grad=GRAD))


def predicate(x: Tensor) -> Tensor:
    return Logic.and_(
        positive_populations(x), limited_population(x), limited_I_from_0(x)
    )


constraint = Constraint(f=lambda x: predicate(x), strength=1.0, gradient_mixer=snr)

generation_options = dict(
    sde_solver=EulerMethod(),
    n_samples=100,
    steps=1000,
    corrector_steps=2,
    final_corrector_steps=2500,
)

CONFIG = ConstrainedGenerationConfig(
    name=NAME + "_inequality",
    constraint=constraint,
    model_path=find(str(MODELS_FOLDER), pattern=f"*{NAME}*"),
    generation_options=generation_options,
    seed=SEED,
)

config/constrained_generation/eSIRS_inequality.py
DavideScassola/score-based-constrained-generation
import fnmatch
import os

import numpy as np
import torch
from matplotlib import image

from src.constrained_generation_config import ConstrainedGenerationConfig
from src.constraints.constraint import Constraint
from src.constraints.gradient_mixer import *
from src.constraints.real_logic import Product as Logic
from src.data import get_MNIST
from src.models.score_based.sdes.sampler import EulerMethod
from src.util import upscale

DEFAULT_MODELS_FOLDER = "artifacts/models"
GRADIENT_APPROXIMATION = linear_interpolation
SEED = 1234
MODEL = "2000"


def mnist_model_path():
    return f"{DEFAULT_MODELS_FOLDER}/{fnmatch.filter(sorted(os.listdir(DEFAULT_MODELS_FOLDER)), f'*mnist*{MODEL}*')[-1]}"


def expand_like(input: torch.Tensor, *, like: torch.Tensor):
    return input.reshape(
        tuple(input.shape) + (1,) * (len(like.shape) - len(input.shape))
    ).expand_as(like)


def to01(x: torch.Tensor):
    x = x - expand_like(x.amin(dim=(1, 2)), like=x)
    return x / expand_like(x.amax(dim=(1, 2)), like=x)


def get_inpainting_constraint_function(
    *,
    image_index: int,
    start_row: int | None = None,
    end_row: int | None = None,
    start_col: int | None = None,
    end_col: int | None = None,
):
    original = get_MNIST(train=False)[image_index]
    rows = slice(start_row, end_row)
    cols = slice(start_col, end_col)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    original_masked = original.detach().clone().unsqueeze(0)
    original_masked_cuda = original.to(device)
    mask = torch.ones_like(original_masked)
    mask[:, rows, cols] = 0.0
    mask_cuda = torch.ones_like(original_masked_cuda)
    mask_cuda[rows, cols] = 0.0
    original_masked *= mask
    original_masked_cuda *= mask_cuda

    image.imsave(f"inpaint_original.png", upscale(original.numpy()), cmap="gray")
    image.imsave(
        f"inpaint_masked.png",
        upscale(np.where(mask[0], original.numpy(), 2.5)),
        cmap="gray",
    )

    def f(x: torch.Tensor) -> torch.Tensor:
        m = original_masked if x.get_device() else original_masked_cuda
        msk = mask if x.get_device() else mask_cuda
        return Logic.equal((m - x) * msk, 0, grad=10).sum(dim=(1, 2))
        # return Logic.and_(Logic.equal((m - x) * msk, 0, grad=1000))

    return f


constraint = Constraint(
    f=get_inpainting_constraint_function(
        image_index=1, start_row=7, end_row=20, start_col=7, end_col=None
    ),
    strength=1,
    gradient_mixer=GRADIENT_APPROXIMATION,
)


generation_options = dict(
    sde_solver=EulerMethod(),
    n_samples=64,
    steps=1000,
    corrector_steps=0,
    final_corrector_steps=0,
)

CONFIG = ConstrainedGenerationConfig(
    name="mnist_inpainting",
    constraint=constraint,
    model_path=mnist_model_path(),
    generation_options=generation_options,
    seed=SEED,
)

config/constrained_generation/mnist_inpainting.py
DavideScassola/score-based-constrained-generation
import fnmatch
import os

import torch
from torch import Tensor

from src.constrained_generation_config import ConstrainedGenerationConfig
from src.constraints.constraint import Constraint
from src.constraints.gradient_mixer import (linear_interpolation,
                                            noise_weighting, snr)
from src.constraints.real_logic import Product as Logic
from src.models.score_based.sdes.sampler import EulerMethod

DEFAULT_MODELS_FOLDER = "artifacts/models"
SCORE_APPROXIMATOR = snr
LOGIC_GRAD = 5.0
SEED = 42
MODEL = "2000"


def mnist_model_path() -> str:
    """
    Picks the last stored model
    """
    return f"{DEFAULT_MODELS_FOLDER}/{fnmatch.filter(sorted(os.listdir(DEFAULT_MODELS_FOLDER)),f'*mnist*{MODEL}*')[-1]}"


def f(x: torch.Tensor) -> Tensor:
    return Logic.equal(x, torch.flip(x, (1,)), grad=LOGIC_GRAD).flatten(1).sum(1)


constraint = Constraint(f=f, strength=1.0, gradient_mixer=SCORE_APPROXIMATOR)

generation_options = dict(
    sde_solver=EulerMethod(),
    n_samples=64,
    steps=1000,
    corrector_steps=1,
    final_corrector_steps=0,
)

CONFIG = ConstrainedGenerationConfig(
    name="mnist_horizontal_symmetry",
    constraint=constraint,
    model_path=mnist_model_path(),
    generation_options=generation_options,
    seed=SEED,
)

config/constrained_generation/mnist_symmetry.py
DavideScassola/score-based-constrained-generation
from src.constants import MODELS_FOLDER
from src.constrained_generation_config import ConstrainedGenerationConfig
from src.constraints.constraint import Constraint
from src.constraints.gradient_mixer import (linear_interpolation,
                                            logistic_interpolation, snr)
from src.constraints.real_logic import Product as Logic
from src.models.score_based.sdes.sampler import EulerMethod
from src.util import find

SEED = 1234

model_path = find(str(MODELS_FOLDER), pattern="*mixture*")

constraint = Constraint(
    f=lambda x: Logic.greater(x, 0.0),
    strength=1.0,
    gradient_mixer=snr,
)

generation_options = dict(
    sde_solver=EulerMethod(),
    steps=1000,
    n_samples=5000,
    corrector_steps=1,
    final_corrector_steps=200,
)

CONFIG = ConstrainedGenerationConfig(
    name="mode_selection",
    constraint=constraint,
    model_path=model_path,
    generation_options=generation_options,
    seed=SEED,
)

config/constrained_generation/mode_selection.py
DavideScassola/score-based-constrained-generation
import fnmatch
import os

import torch

from src.constrained_generation_config import ConstrainedGenerationConfig
from src.constraints.constraint import Constraint
from src.constraints.gradient_mixer import (linear_interpolation,
                                            noise_weighting, snr)
from src.constraints.real_logic import Product as Logic
from src.models.score_based.sdes.sampler import EulerMethod

DEFAULT_MODELS_FOLDER = "artifacts/models"
GAP = 0
FILL_DIFFERENCE = 1
LOGIC_K = 80
STRENGTH = 1
GRADIENT_APPROXIMATION = snr
SEED = 1234
MODEL = "2000"


def mnist_model_path():
    """
    Picks the last stored model
    """
    return f"{DEFAULT_MODELS_FOLDER}/{fnmatch.filter(sorted(os.listdir(DEFAULT_MODELS_FOLDER)), f'*mnist*{MODEL}*')[-1]}"


""""
def expand_like(input: torch.Tensor, *, like: torch.Tensor):
    return input.reshape(
        tuple(input.shape) + (1,) * (len(like.shape) - len(input.shape))
    ).expand_as(like)

def to01(x: torch.Tensor):
    x = x - expand_like(x.amin(dim=(1, 2)), like=x)
    return x / expand_like(x.amax(dim=(1, 2)), like=x)

def f(x: torch.Tensor):
    return Logic.greater(to01(x).mean(dim=(1,2)), 1.0)
"""


def relative_filling_down(x: torch.Tensor):
    return Logic.smaller(
        x[:, : 14 - GAP, :].mean(dim=(1, 2)) + FILL_DIFFERENCE,
        x[:, 14 + GAP :, :].mean(dim=(1, 2)),
        grad=LOGIC_K,
    )


constraint = Constraint(
    f=relative_filling_down, strength=STRENGTH, gradient_mixer=GRADIENT_APPROXIMATION
)

generation_options = dict(
    sde_solver=EulerMethod(),
    n_samples=64,
    steps=1000,
    corrector_steps=2,
    final_corrector_steps=0,
)

CONFIG = ConstrainedGenerationConfig(
    name="mnist_relative_filling_down",
    constraint=constraint,
    model_path=mnist_model_path(),
    generation_options=generation_options,
    seed=SEED,
)

config/constrained_generation/relative_filling_down.py
DavideScassola/score-based-constrained-generation
import fnmatch
import os

import torch

from src.constrained_generation_config import ConstrainedGenerationConfig
from src.constraints.constraint import Constraint
from src.constraints.gradient_mixer import noise_weighting, snr
from src.constraints.real_logic import Product as Logic
from src.models.score_based.sdes.sampler import EulerMethod

DEFAULT_MODELS_FOLDER = "artifacts/models"
GAP = 0
FILL_DIFFERENCE = 1
LOGIC_K = 80
STRENGTH = 1
GRADIENT_APPROXIMATION = snr
SEED = 1234
MODEL = "2000"


def mnist_model_path():
    """
    Picks the last stored model
    """
    return f"{DEFAULT_MODELS_FOLDER}/{fnmatch.filter(sorted(os.listdir(DEFAULT_MODELS_FOLDER)), '*mnist*2000*')[-1]}"


""""
def expand_like(input: torch.Tensor, *, like: torch.Tensor):
    return input.reshape(
        tuple(input.shape) + (1,) * (len(like.shape) - len(input.shape))
    ).expand_as(like)

def to01(x: torch.Tensor):
    x = x - expand_like(x.amin(dim=(1, 2)), like=x)
    return x / expand_like(x.amax(dim=(1, 2)), like=x)

def f(x: torch.Tensor):
    return Logic.greater(to01(x).mean(dim=(1,2)), 1.0)
"""


def relative_filling_down(x: torch.Tensor):
    return Logic.greater(
        x[:, : 14 - GAP, :].mean(dim=(1, 2)) - FILL_DIFFERENCE,
        x[:, 14 + GAP :, :].mean(dim=(1, 2)),
    )


constraint = Constraint(
    f=relative_filling_down, strength=STRENGTH, gradient_mixer=GRADIENT_APPROXIMATION
)

generation_options = dict(
    sde_solver=EulerMethod(),
    n_samples=64,
    steps=1000,
    corrector_steps=2,
    final_corrector_steps=0,
)

CONFIG = ConstrainedGenerationConfig(
    name="mnist_relative_filling_up",
    constraint=constraint,
    model_path=mnist_model_path(),
    generation_options=generation_options,
    seed=SEED,
)

config/constrained_generation/relative_filling_up.py
DavideScassola/score-based-constrained-generation
import os

from src.constants import MODELS_FOLDER
from src.constrained_generation_config import ConstrainedGenerationConfig
from src.constraints.constraint import Constraint
from src.constraints.gradient_mixer import (linear_interpolation,
                                            noise_weighting, snr)
from src.constraints.real_logic import Product as Logic
from src.models.score_based.sdes.sampler import EulerMethod
from src.util import find, names_index_map

model_path = find(str(MODELS_FOLDER), pattern="*white_wine*")

i = names_index_map("data/wine_quality_white.csv")

alcohol = i["alcohol"]
fixed_acidity = i["fixed acidity"]
residual_sugar = i["residual sugar"]
citric_acid = i["citric acid"]


GRAD = 50.0
SCORE_APPROXIMATOR = linear_interpolation
SEED = 1234

predicate = lambda x: Logic.and_(
    Logic.or_(
        Logic.in_(x[:, fixed_acidity], 5.0, 6.0, grad=GRAD),
        Logic.in_(x[:, fixed_acidity], 8.0, 9.0, grad=GRAD),
    ),
    Logic.greater(x[:, alcohol], x2=11.0, grad=GRAD),
    Logic.or_(
        Logic.greater(x[:, residual_sugar], 5.0, grad=GRAD),
        Logic.greater(x[:, citric_acid], 0.5, grad=GRAD),
    ),
)


constraint = Constraint(
    f=lambda x: predicate(x), strength=1.0, gradient_mixer=SCORE_APPROXIMATOR
)

generation_options = dict(
    sde_solver=EulerMethod(),
    n_samples=5000,
    steps=1000,
    corrector_steps=1,
    final_corrector_steps=4000,
)

CONFIG = ConstrainedGenerationConfig(
    name="white_wine",
    constraint=constraint,
    model_path=model_path,
    generation_options=generation_options,
    seed=SEED,
)

config/constrained_generation/white_wine.py
DavideScassola/score-based-constrained-generation
from functools import partial

import torch

from src.data import Dataset
from src.models.score_based.nn.optimization import Optimization
from src.models.score_based.score_function import MLP
from src.models.score_based.score_matching import (
    accumulated_denoising_score_matching_loss, denoising_score_matching_loss,
    hybrid_ssm_dsm_loss, sliced_score_matching_loss)
from src.models.score_based.sdes.sampler import EulerMethod
from src.models.score_based.sdes.sde import VE, subVP
from src.models.score_based.tabular_score_based_sde import TabularScoreBasedSde
from src.preprocessors.mean_std_normalizer import MeanStdNormalizer
from src.train_config import TrainConfig

dataset = Dataset(path="data/airbnb.csv", train_proportion=1.0)

model = TabularScoreBasedSde(
    tensor_preprocessors=[
        MeanStdNormalizer(),
    ],
    sde=subVP(),
    score_function_class=MLP,
    score_function_hyperparameters={
        "hidden_channels": (100, 50),
        "activation_layer": torch.nn.SiLU,
        "add_prior_score_bias": True,
        "rescale_factor_limit": 4,
    },
    optimization=Optimization(
        epochs=100,
        batch_size=1000,
        optimizer_class=torch.optim.RAdam,
        optimizer_hyperparameters={"lr": 1e-3},
        # parameters_momentum=0.999,
    ),
    score_matching_loss=partial(hybrid_ssm_dsm_loss, w=0.1),
)

generation_options = dict(
    sde_solver=EulerMethod(),
    n_samples=3000,
    steps=1000,
    corrector_steps=1,
    final_corrector_steps=100,
)

CONFIG = TrainConfig(
    name="airbnb",
    dataset=dataset,
    model=model,
    generation_options=generation_options,
)

config/train/airbnb.py
DavideScassola/score-based-constrained-generation
import torch

from src.data import Dataset
from src.models.score_based.image_score_based_sde import ImageScoreBasedSde
from src.models.score_based.nn.optimization import Optimization
from src.models.score_based.score_function import MLP, Unet
from src.models.score_based.sdes.sampler import EulerMethod
from src.models.score_based.sdes.sde import VE, subVP
from src.preprocessors.mean_std_normalizer import MeanStdNormalizer
from src.preprocessors.rescaler import Rescaler
from src.train_config import TrainConfig

dataset = Dataset(path="MNIST", train_proportion=0.8)

model = ImageScoreBasedSde(
    tensor_preprocessors=[],
    sde=VE(),
    score_function_class=Unet,
    score_function_hyperparameters={},
    optimization=Optimization(
        epochs=1500,
        batch_size=256,
        optimizer_class=torch.optim.RAdam,
        optimizer_hyperparameters={"lr": 1e-3},
    ),
)

generation_options = dict(
    sde_solver=EulerMethod(), n_samples=20, steps=1000, corrector_steps=0
)

CONFIG = TrainConfig(
    name="mnist", dataset=dataset, model=model, generation_options=generation_options
)

config/train/debug.py
DavideScassola/score-based-constrained-generation
from functools import partial

import torch

from src.data import Dataset
from src.models.score_based.nn.optimization import Optimization
from src.models.score_based.score_function import MLP
from src.models.score_based.score_matching import \
    accumulated_denoising_score_matching_loss
from src.models.score_based.sdes.sampler import EulerMethod
from src.models.score_based.sdes.sde import VE, subVP
from src.models.score_based.time_series_score_based_sde import \
    TimeSeriesScoreBasedSde
from src.preprocessors.mean_std_normalizer import MeanStdNormalizer
from src.train_config import TrainConfig

dataset = Dataset(path="data/eSIRS.npy", train_proportion=1.0)

model = TimeSeriesScoreBasedSde(
    tensor_preprocessors=[MeanStdNormalizer()],
    sde=subVP(),
    score_function_class=MLP,
    score_function_hyperparameters={
        "hidden_channels": (200, 100, 100),
        "activation_layer": torch.nn.SiLU,
        "add_prior_score_bias": True,
        "rescale_factor_limit": 1,
    },
    optimization=Optimization(
        epochs=2000,
        batch_size=1000,
        optimizer_class=torch.optim.RAdam,
        optimizer_hyperparameters={"lr": 5e-3},
    ),
    score_matching_loss=partial(accumulated_denoising_score_matching_loss, n=10),
    names=["S", "I"],
)

generation_options = dict(
    sde_solver=EulerMethod(), steps=1000, corrector_steps=0, n_samples=100
)

CONFIG = TrainConfig(
    name="eSIRS", dataset=dataset, model=model, generation_options=generation_options
)

config/train/eSIRS.py
DavideScassola/score-based-constrained-generation
import torch

from src.data import Dataset
from src.models.score_based.image_score_based_sde import ImageScoreBasedSde
from src.models.score_based.nn.optimization import Optimization
from src.models.score_based.score_function import MLP, Unet
from src.models.score_based.sdes.sampler import EulerMethod
from src.models.score_based.sdes.sde import VE, subVP
from src.preprocessors.mean_std_normalizer import MeanStdNormalizer
from src.preprocessors.rescaler import Rescaler
from src.train_config import TrainConfig

dataset = Dataset(path="MNIST", train_proportion=0.8)

model = ImageScoreBasedSde(
    tensor_preprocessors=[],
    sde=VE(),
    score_function_class=Unet,
    score_function_hyperparameters={},
    optimization=Optimization(
        epochs=250,
        batch_size=256,
        optimizer_class=torch.optim.RAdam,
        optimizer_hyperparameters={"lr": 1e-3},
    ),
)

generation_options = dict(
    sde_solver=EulerMethod(), n_samples=20, steps=1000, corrector_steps=0
)

CONFIG = TrainConfig(
    name="mnist", dataset=dataset, model=model, generation_options=generation_options
)

config/train/mnist.py
DavideScassola/score-based-constrained-generation
from functools import partial

import torch

from src.data import Dataset
from src.models.score_based.nn.optimization import Optimization
from src.models.score_based.score_function import MLP, GineMLP
from src.models.score_based.score_matching import (
    accumulated_denoising_score_matching_loss, hybrid_ssm_dsm_loss,
    sliced_score_matching_loss)
from src.models.score_based.sdes.sampler import EulerMethod, TimeAdaptiveEuler
from src.models.score_based.sdes.sde import VE, subVP
from src.models.score_based.tabular_score_based_sde import TabularScoreBasedSde
from src.preprocessors.mean_std_normalizer import MeanStdNormalizer
from src.train_config import TrainConfig

dataset = Dataset(path="data/mono_mixture.csv", train_proportion=0.8)

model = TabularScoreBasedSde(
    tensor_preprocessors=[
        MeanStdNormalizer(),
    ],
    sde=subVP(),
    score_function_class=MLP,
    score_function_hyperparameters={
        "hidden_channels": (64, 64),
        "activation_layer": torch.nn.SiLU,
        "add_prior_score_bias": True,
        "rescale_factor_limit": 5,
    },
    optimization=Optimization(
        epochs=600,
        batch_size=500,
        optimizer_class=torch.optim.AdamW,
        optimizer_hyperparameters={"lr": 1e-2},
        # parameters_momentum=0.99,
    ),
    score_matching_loss=partial(
        hybrid_ssm_dsm_loss,
        w=0.01,
    ),
)

generation_options = dict(
    sde_solver=EulerMethod(),
    steps=1000,
    n_samples=2000,
    corrector_steps=1,
    final_corrector_steps=200,
)

CONFIG = TrainConfig(
    name="mixture", dataset=dataset, model=model, generation_options=generation_options
)

config/train/mono_mixture.py
DavideScassola/score-based-constrained-generation
from functools import partial

import torch

from src.data import Dataset
from src.models.score_based.nn.optimization import Optimization
from src.models.score_based.score_function import MLP
from src.models.score_based.score_matching import \
    accumulated_denoising_score_matching_loss
from src.models.score_based.sdes.sampler import EulerMethod, TimeAdaptiveEuler
from src.models.score_based.sdes.sde import VE, subVP
from src.models.score_based.tabular_score_based_sde import TabularScoreBasedSde
from src.preprocessors.mean_std_normalizer import MeanStdNormalizer
from src.train_config import TrainConfig

dataset = Dataset(path="data/mono_normal.csv", train_proportion=0.8)

model = TabularScoreBasedSde(
    tensor_preprocessors=[
        MeanStdNormalizer(),
    ],
    sde=VE(),
    score_function_class=MLP,
    score_function_hyperparameters={
        "hidden_channels": (32,),
        "activation_layer": torch.nn.SiLU,
        "add_prior_score_bias": True,
        "rescale_factor_limit": 1,
    },
    optimization=Optimization(
        epochs=300,
        batch_size=200,
        optimizer_class=torch.optim.RAdam,
        optimizer_hyperparameters={"lr": 2e-2},
    ),
    score_matching_loss=partial(
        accumulated_denoising_score_matching_loss, reciprocal_distribution_t=True, n=5
    ),
)

generation_options = dict(
    sde_solver=EulerMethod(),
    steps=1000,
    n_samples=1000,
    corrector_steps=0,
    final_corrector_steps=1000,
)

CONFIG = TrainConfig(
    name="normal", dataset=dataset, model=model, generation_options=generation_options
)

config/train/mono_normal.py
DavideScassola/score-based-constrained-generation
from src.data import Dataset
from src.models.multivariate_gaussian import MultivariateGaussian
from src.preprocessors.quantile_normalizer import QuantileNormalizer
from src.preprocessors.to_numpy import ToNumpy
from src.train_config import TrainConfig

dataset = Dataset(path="data/uniform_exp_mixture.csv", train_proportion=0.8)

model = MultivariateGaussian(preprocessors=([ToNumpy(), QuantileNormalizer()]))

generation_options = dict(n_samples=1000)

CONFIG = TrainConfig(
    name="multivariate_gaussian",
    dataset=dataset,
    model=model,
    generation_options=generation_options,
)

config/train/multivariate_gaussian.py
DavideScassola/score-based-constrained-generation
from functools import partial

import torch

from src.data import Dataset
from src.models.score_based.nn.optimization import Optimization
from src.models.score_based.score_function import MLP
from src.models.score_based.score_matching import (
    accumulated_denoising_score_matching_loss, hybrid_ssm_dsm_loss,
    sliced_score_matching_loss)
from src.models.score_based.sdes.sampler import EulerMethod
from src.models.score_based.sdes.sde import VE, subVP
from src.models.score_based.tabular_score_based_sde import TabularScoreBasedSde
from src.preprocessors.mean_std_normalizer import MeanStdNormalizer
from src.train_config import TrainConfig

dataset = Dataset(path="data/wine_quality_white.csv", train_proportion=0.8)

model = TabularScoreBasedSde(
    tensor_preprocessors=[
        MeanStdNormalizer(),
    ],
    sde=subVP(),
    score_function_class=MLP,
    score_function_hyperparameters={
        "hidden_channels": (150, 100, 50),
        "activation_layer": torch.nn.SiLU,
        "add_prior_score_bias": True,
        "rescale_factor_limit": 1,
    },
    optimization=Optimization(
        epochs=3000,
        batch_size=500,
        optimizer_class=torch.optim.RAdam,
        optimizer_hyperparameters={"lr": 1e-3},
        # parameters_momentum=0.999,
    ),
    score_matching_loss=partial(
        accumulated_denoising_score_matching_loss, n=10, reciprocal_distribution_t=True
    ),
)

generation_options = dict(
    sde_solver=EulerMethod(),
    n_samples=1000,
    steps=1000,
    corrector_steps=1,
    final_corrector_steps=1000,
)

CONFIG = TrainConfig(
    name="white_wine",
    dataset=dataset,
    model=model,
    generation_options=generation_options,
)

config/train/white_wine.py
DavideScassola/score-based-constrained-generation
import os
import urllib.request
from pathlib import Path

import numpy as np
import pandas as pd

CSV_URL = "https://datahub.io/machine-learning/adult/r/adult.csv"


def self_base_name() -> str:
    return Path(os.path.basename(__file__)).stem


def csv_path() -> str:
    return f"data/{self_base_name()}.csv"


def main():
    os.makedirs("data", exist_ok=True)
    urllib.request.urlretrieve(CSV_URL, csv_path())


if __name__ == "__main__":
    main()

datasets_scripts/adult.py
DavideScassola/score-based-constrained-generation
import os
import urllib.request
import zipfile
from pathlib import Path

import pandas as pd

ZIP_URL = "https://www.dropbox.com/scl/fi/wyz56vpff026gou08s230/airbnb_ny.zip?rlkey=bix104azzv9uc9uqs9yskh3po&dl=1"
CSV_PATH = "data/AB_NYC_2019.csv"
ZIP_PATH = "data/airbnb_ny.zip"


def self_base_name() -> str:
    return Path(os.path.basename(__file__)).stem


def csv_path() -> str:
    return f"data/{self_base_name()}.csv"


def main():
    urllib.request.urlretrieve(ZIP_URL, ZIP_PATH)

    with zipfile.ZipFile(ZIP_PATH, "r") as zip_ref:
        # zip_ref.extractall(directory_to_extract_to)
        zip_ref.extract("AB_NYC_2019.csv", "data")

    pd.read_csv(CSV_PATH)[["latitude", "longitude"]].to_csv(csv_path(), index=False)
    os.remove(ZIP_PATH)
    os.remove(CSV_PATH)


if __name__ == "__main__":
    main()

datasets_scripts/airbnb.py
DavideScassola/score-based-constrained-generation
import gillespy2
import numpy as np
import tqdm


def sample_random_partition(*, total: int, partitions: int) -> np.ndarray:
    x = np.random.dirichlet(alpha=np.ones(partitions), size=1).squeeze()
    x = np.round(x * total)
    x[-1] = total - np.sum(x[:-1])
    return x


class eSIRS(gillespy2.Model):
    def __init__(self, total_population, endtime, n_steps):
        # First call the gillespy2.Model initializer.
        super().__init__("eSIRS")
        # Define parameters for the rates of creation and dissociation.
        k1 = gillespy2.Parameter(name="k1", expression=2.36012158 / total_population)
        k2 = gillespy2.Parameter(name="k2", expression=1.6711464)
        k3 = gillespy2.Parameter(name="k3", expression=0.90665231)
        k4 = gillespy2.Parameter(name="k4", expression=0.63583386)

        self.add_parameter([k1, k2, k3, k4])
        self.total_population = total_population
        # Define variables for the molecular species representing M & D.
        S = gillespy2.Species(name="S", initial_value=total_population)
        I = gillespy2.Species(name="I", initial_value=0)
        R = gillespy2.Species(name="R", initial_value=0)
        self.species_list = [S, I, R]
        self.add_species(self.species_list)

        # The list of reactants and products for a Reaction object are
        # each a Python dictionary in which the dictionary keys are
        # Species objects and the values are stoichiometries of the
        # species in the reaction.

        # infection
        r1 = gillespy2.Reaction(
            name="r1", rate=k1, reactants={S: 1, I: 1}, products={I: 2}
        )
        # recovery
        r2 = gillespy2.Reaction(name="r2", rate=k2, reactants={I: 1}, products={R: 1})

        # loss of immunity
        r3 = gillespy2.Reaction(name="r3", rate=k3, reactants={R: 1}, products={S: 1})

        # external infection
        r4 = gillespy2.Reaction(name="r4", rate=k4, reactants={S: 1}, products={I: 1})

        self.add_reaction([r1, r2, r3, r4])

        # Set the timespan for the simulation.
        self.timespan(np.linspace(0, endtime, n_steps))

    def set_initial_state(self) -> None:
        x = sample_random_partition(
            total=self.total_population, partitions=len(self.species_list)
        )
        for i in range(len(self.species_list)):
            self.species_list[i].initial_value = x[i]


if __name__ == "__main__":
    N_SAMPLES = 10_000
    STARTING_POINTS = 1000  # ideally it would be the same as the number of samples, but it's not efficient (probably there is a solution)

    total_population = 100
    n_steps = 30
    endtime = 3.2
    trajectories_per_sarting_point = N_SAMPLES // STARTING_POINTS

    model = eSIRS(total_population=total_population, n_steps=n_steps, endtime=endtime)

    trajectories = np.empty(
        (
            STARTING_POINTS,
            trajectories_per_sarting_point,
            n_steps,
            len(model.species_list) - 1,
        )
    )

    for i in tqdm.tqdm(range(STARTING_POINTS)):
        model.set_initial_state()
        x = model.run(number_of_trajectories=trajectories_per_sarting_point)
        for j in range(trajectories_per_sarting_point):
            trajectories[i, j, :, 0] = x[j]["S"]
            trajectories[i, j, :, 1] = x[j]["I"]

    X = trajectories.reshape(
        (STARTING_POINTS * trajectories_per_sarting_point, n_steps, -1)
    )

    np.random.shuffle(X)

    filename = f"data/eSIRS.npy"

    np.save(filename, arr=X, allow_pickle=False)

datasets_scripts/eSIRS.py
DavideScassola/score-based-constrained-generation
import os
import urllib.request
import zipfile
from pathlib import Path

import numpy as np
import pandas as pd

ZIP_URL = "https://archive.ics.uci.edu/static/public/235/individual+household+electric+power+consumption.zip"
FILE_NAME = "data/household_power_consumption.txt"


def self_base_name() -> str:
    return Path(os.path.basename(__file__)).stem


def zip_path() -> str:
    return f"data/{self_base_name()}.zip"


def csv_path() -> str:
    return f"data/{self_base_name()}.csv"


def extract_zip(archive: str) -> None:
    with zipfile.ZipFile(file=archive, mode="r") as zip_ref:
        zip_ref.extractall(path=Path(archive).parent)
    os.remove(path=archive)


def main():
    os.makedirs(name="data", exist_ok=True)
    if not os.path.isfile(csv_path()):
        urllib.request.urlretrieve(url=ZIP_URL, filename=zip_path())
        extract_zip(archive=zip_path())
        os.rename(src=FILE_NAME, dst=csv_path())
    else:
        print(f"{csv_path()} already present")

    df = pd.read_csv(
        csv_path(),
        sep=";",
        parse_dates={"dt": ["Date", "Time"]},
        dayfirst=True,
        low_memory=False,
        na_values=["nan", "?"],
        index_col="dt",
    )

    df = df.resample("H").median()  # first ?
    df = df[
        "2006-12-17":"2010-11-25"
    ]  # .dropna() # discard first and last day since they are not complete
    df_grouped = df.groupby(df.index.to_period("D"))
    X = np.stack(df_grouped.apply(np.array))
    X = X[~np.any(np.isnan(X), axis=(1, 2))]


if __name__ == "__main__":
    main()

datasets_scripts/household_power_consumption.py
DavideScassola/score-based-constrained-generation
import os
from pathlib import Path

import numpy as np
import pandas as pd

CLASSES = {"A": 0.1, "B": 0.45, "C": 0.3, "D": 0.15}


def self_base_name() -> str:
    return Path(os.path.basename(__file__)).stem


def csv_path() -> str:
    return f"data/{self_base_name()}.csv"


def generate(size: int = 1000) -> pd.DataFrame:
    c = np.random.choice(list(CLASSES.keys()), size=size, p=list(CLASSES.values()))
    return pd.DataFrame({f"categorical({str(CLASSES)})": c})


def store(df: pd.DataFrame):
    os.makedirs("data", exist_ok=True)
    generate(1000).to_csv(csv_path(), index=False)


def main():
    store(generate())


if __name__ == "__main__":
    main()

datasets_scripts/mono_categorical.py
DavideScassola/score-based-constrained-generation
import os

import numpy as np
import pandas as pd


def generate(size: int = 1000) -> pd.DataFrame:
    return pd.DataFrame({"exponential(1)": np.random.exponential(1, size)})


def store(df: pd.DataFrame):
    os.makedirs("data", exist_ok=True)
    name = os.path.basename(__file__).split(".py")[0]
    df.to_csv(f"data/{name}.csv", index=False)


if __name__ == "__main__":
    store(generate())

datasets_scripts/mono_exponential.py
DavideScassola/score-based-constrained-generation
import os

import numpy as np
import pandas as pd


def random_mixture(mix_and_samplers: tuple, size: int):
    mix_coefficients = [k["p"] for k in mix_and_samplers]
    samplers = [k["sampler"] for k in mix_and_samplers]

    z = np.random.multinomial(n=size, pvals=mix_coefficients)
    x = np.concatenate([samplers[i](n) for i, n in enumerate(z)], axis=0)
    np.random.shuffle(x)
    return x


def generate(size: int = 1000) -> pd.DataFrame:
    return pd.DataFrame(
        {
            "mixture(0.5: N(-3, 0.5), 0.5: N(4, 1))": random_mixture(
                (
                    {"p": 0.5, "sampler": lambda s: np.random.normal(-3, 0.5, s)},
                    {"p": 0.5, "sampler": lambda s: np.random.normal(4, 1, s)},
                ),
                size,
            ),
        }
    )


def store(df: pd.DataFrame):
    name = os.path.basename(__file__).split(".py")[0]
    os.makedirs("data", exist_ok=True)
    df.to_csv(f"data/{name}.csv", index=False)


if __name__ == "__main__":
    store(generate())

datasets_scripts/mono_mixture.py
DavideScassola/score-based-constrained-generation
import os

import numpy as np
import pandas as pd


def generate(size: int = 1000) -> pd.DataFrame:
    return pd.DataFrame({"N(0, 1)": np.random.normal(0, 1, size)})


def store(df: pd.DataFrame):
    os.makedirs("data", exist_ok=True)
    name = os.path.basename(__file__).split(".py")[0]
    df.to_csv(f"data/{name}.csv", index=False)


if __name__ == "__main__":
    store(generate())

datasets_scripts/mono_normal.py
DavideScassola/score-based-constrained-generation
import os

import numpy as np
import pandas as pd


def generate(size: int = 1000) -> pd.DataFrame:
    return pd.DataFrame({"uniform(0, 1)": np.random.uniform(0, 1, size)})


def store(df: pd.DataFrame):
    os.makedirs("data", exist_ok=True)
    name = os.path.basename(__file__).split(".py")[0]
    df.to_csv(f"data/{name}.csv", index=False)


if __name__ == "__main__":
    store(generate())

datasets_scripts/mono_uniform.py
DavideScassola/score-based-constrained-generation
import os

import numpy as np
import pandas as pd


def generate(size: int = 1000) -> pd.DataFrame:
    cov = [[2.0, 1.0, -1.0], [1.0, 1.0, 0.2], [-1, 0.2, 1.5]]
    mean = [0, -1, 2]
    x = np.random.multivariate_normal(mean=mean, cov=cov, size=size)
    return pd.DataFrame(x)


def store(df: pd.DataFrame):
    os.makedirs("data", exist_ok=True)
    name = os.path.basename(__file__).split(".py")[0]
    df.to_csv(f"data/{name}.csv", index=False)


if __name__ == "__main__":
    store(generate())

datasets_scripts/multivariate_gaussian.py
DavideScassola/score-based-constrained-generation
import os

import numpy as np
import pandas as pd


def random_mixture(mix_and_samplers: tuple, size: int):
    mix_coefficients = [k["p"] for k in mix_and_samplers]
    samplers = [k["sampler"] for k in mix_and_samplers]

    z = np.random.multinomial(n=size, pvals=mix_coefficients)
    x = np.concatenate([samplers[i](n) for i, n in enumerate(z)], axis=0)
    np.random.shuffle(x)
    return x


def generate(size: int = 1000) -> pd.DataFrame:
    return pd.DataFrame(
        {
            "uniform(-120, 20)": np.random.uniform(-120, -20, size),
            "exponential(0.1)": np.random.exponential(0.1, size),
            "mixture(0.2: N(-3, 0.2), 0.8: N(500, 1))": random_mixture(
                (
                    {"p": 0.2, "sampler": lambda s: np.random.normal(-3, 0.2, s)},
                    {"p": 0.8, "sampler": lambda s: np.random.normal(500, 1, s)},
                ),
                size,
            ),
        }
    )


def store(df: pd.DataFrame):
    os.makedirs("data", exist_ok=True)
    name = os.path.basename(__file__).split(".py")[0]
    df.to_csv(f"data/{name}.csv", index=False)


if __name__ == "__main__":
    store(generate())

datasets_scripts/uniform_exp_mixture.py
DavideScassola/score-based-constrained-generation
import os
import urllib.request
from pathlib import Path

import numpy as np
import pandas as pd

CSV_URL = "https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-white.csv"


def self_base_name() -> str:
    return Path(os.path.basename(__file__)).stem


def csv_path() -> str:
    return f"data/{self_base_name()}.csv"


def main():
    os.makedirs("data", exist_ok=True)
    urllib.request.urlretrieve(CSV_URL, csv_path())
    df = pd.read_csv(csv_path(), sep=";")
    df.drop("quality", axis=1, inplace=True)
    df.to_csv(csv_path(), index=False)


if __name__ == "__main__":
    main()

datasets_scripts/wine_quality_white.py
DavideScassola/score-based-constrained-generation
import argparse

from src.constrained_generation_controller import run_constrained_generation

DEBUG_CONFIG = "config/constrained_generation/debug.py"


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("config_file", type=str, nargs="?", default=DEBUG_CONFIG)
    return parser.parse_args()


if __name__ == "__main__":
    run_constrained_generation(parse_args().config_file)

main-generate.py
DavideScassola/score-based-constrained-generation
import argparse

from src.train_controller import run_model_training

DEBUG_CONFIG = "config/train/debug.py"


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("config_file", type=str, nargs="?", default=DEBUG_CONFIG)
    return parser.parse_args()


if __name__ == "__main__":
    run_model_training(parse_args().config_file)

main-train.py
DavideScassola/score-based-constrained-generation
import gillespy2
import numpy as np
import tqdm


def sample_random_partition(*, total: int, partitions: int) -> np.ndarray:
    x = np.random.dirichlet(alpha=np.ones(partitions), size=1).squeeze()
    x = np.round(x * total)
    x[-1] = total - np.sum(x[:-1])
    return x


S = 80
I = 10


class eSIRS(gillespy2.Model):
    def __init__(self, total_population, endtime, n_steps):
        # First call the gillespy2.Model initializer.
        super().__init__("eSIRS")
        # Define parameters for the rates of creation and dissociation.
        k1 = gillespy2.Parameter(name="k1", expression=2.36012158 / total_population)
        k2 = gillespy2.Parameter(name="k2", expression=1.6711464)
        k3 = gillespy2.Parameter(name="k3", expression=0.90665231)
        k4 = gillespy2.Parameter(name="k4", expression=0.63583386)

        self.add_parameter([k1, k2, k3, k4])
        self.total_population = total_population
        # Define variables for the molecular species representing M & D.
        S = gillespy2.Species(name="S", initial_value=total_population)
        I = gillespy2.Species(name="I", initial_value=0)
        R = gillespy2.Species(name="R", initial_value=0)
        self.species_list = [S, I, R]
        self.add_species(self.species_list)

        # The list of reactants and products for a Reaction object are
        # each a Python dictionary in which the dictionary keys are
        # Species objects and the values are stoichiometries of the
        # species in the reaction.

        # infection
        r1 = gillespy2.Reaction(
            name="r1", rate=k1, reactants={S: 1, I: 1}, products={I: 2}
        )
        # recovery
        r2 = gillespy2.Reaction(name="r2", rate=k2, reactants={I: 1}, products={R: 1})

        # loss of immunity
        r3 = gillespy2.Reaction(name="r3", rate=k3, reactants={R: 1}, products={S: 1})

        # external infection
        r4 = gillespy2.Reaction(name="r4", rate=k4, reactants={S: 1}, products={I: 1})

        self.add_reaction([r1, r2, r3, r4])

        # Set the timespan for the simulation.
        self.timespan(np.linspace(0, endtime, n_steps))

    def set_initial_state(self) -> None:
        x = (S, I, self.total_population - S - I)
        for i in range(len(self.species_list)):
            self.species_list[i].initial_value = x[i]


def simulate_bridging():
    print("simulating trajectories...")

    N_SAMPLES = 1_000_000
    STARTING_POINTS = 1

    total_population = 100
    n_steps = 30
    endtime = 3.2
    trajectories_per_sarting_point = N_SAMPLES // STARTING_POINTS

    model = eSIRS(total_population=total_population, n_steps=n_steps, endtime=endtime)

    trajectories = np.empty(
        (
            STARTING_POINTS,
            trajectories_per_sarting_point,
            n_steps,
            len(model.species_list) - 1,
        )
    )

    for i in tqdm.tqdm(range(STARTING_POINTS)):
        model.set_initial_state()
        x = model.run(number_of_trajectories=trajectories_per_sarting_point)
        for j in range(trajectories_per_sarting_point):
            trajectories[i, j, :, 0] = x[j]["S"]
            trajectories[i, j, :, 1] = x[j]["I"]

    X = trajectories.reshape(
        (STARTING_POINTS * trajectories_per_sarting_point, n_steps, -1)
    )

    np.random.shuffle(X)

    filename = f"paper_utils/eSIRS_S0={S}_I0={I}_S25=30.npy"

    X = X[X[:, 25, 0] == 30]

    np.save(filename, arr=X, allow_pickle=False)


if __name__ == "__main__":
    simulate_bridging()

paper_utils/eSIRS_bridging_simulate.py
DavideScassola/score-based-constrained-generation
import json
import random
import sys

import matplotlib.pyplot as plt
import numpy as np
import torch

BASE = 3.9
SEED = 1234

_, axes = plt.subplots(nrows=2, ncols=1, figsize=(BASE, BASE / 1.2), sharex=True)


def set_seeds(seed: int):
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)


def store_json(d: dict, *, file: str) -> None:
    with open(file, "w") as f:
        json.dump(d, f, indent=4)


def positivity(x: np.ndarray):
    return np.all((x >= 0).reshape(len(x), -1), axis=-1)


def constant_population(x: np.ndarray):
    return np.all(np.sum(x, axis=-1) <= 100, axis=-1)


def less_20_infected(x: np.ndarray):
    return np.all(x[:, :, 1] <= 20, axis=-1)


def time_series_plot(
    x: np.ndarray, *, features_names: list, color: str, label: str, alpha=0.5
) -> None:
    k = len(features_names)

    for i, name in enumerate(features_names):
        axes[i].plot(x[:, :, i].T, alpha=alpha, color=color, linewidth=0.8)
        axes[i].set_ylabel(name)
    plt.xlabel("time")


set_seeds(SEED)

path = sys.argv[1]
# config_module = load_module(path=path + '/config.py')


colors = {"Guidance": "orange", "No guidance": "blue"}
x = {
    "Guidance": np.load(path + "/report/guidance.npy"),
    "No guidance": np.load(path + "/report/no_guidance.npy"),
}

x["Guidance"] = np.round(x["Guidance"][np.random.permutation(len(x["Guidance"]))])
x["No guidance"] = np.round(
    x["No guidance"][np.random.permutation(len(x["No guidance"]))]
)

index = {"S": 0, "I": 1}

ALPHA = 0.6
N_SAMPLES = 10
bridging = np.mean(np.round(x["Guidance"][:, 0, 0]) == 95.0) > 0.5

stats = {}

for i in ("No guidance", "Guidance"):
    samples = np.round(x[i])
    sub_sample = samples[:N_SAMPLES]
    time_series_plot(
        sub_sample, features_names=["S", "I"], color=colors[i], alpha=ALPHA, label=i
    )
    stats[f"{i} positivity"] = np.mean(positivity(samples))
    stats[f"{i} constant population: "] = np.mean(constant_population(samples))

for ax in axes:
    ax.grid(alpha=0.4, zorder=-1)
    ax.set_ylim(0, None)

# print("less 20 infected: ", np.mean(less_20_infected(x["Guidance"])))
stats["less 20 infected"] = np.mean(less_20_infected(x["Guidance"]))
axes[1].axhline(y=20, color="r", linestyle="--", alpha=0.6, linewidth=1.0)


legend = axes[1].legend(("No guidance", "Guidance"))
legend.legend_handles[0].set_color("blue")
legend.legend_handles[1].set_color("orange")
print(stats)

plt.tight_layout(pad=0.1)
plt.savefig(path + "/paper_plot.svg")
store_json(stats, file=path + "/plot_stats.json")
plt.show()

paper_utils/eSIRS_plot.py
DavideScassola/score-based-constrained-generation
import json
import os
import random
import sys

import matplotlib.pyplot as plt
import numpy as np
import torch
from eSIRS_bridging_simulate import simulate_bridging

BASE = 3.9
SEED = 1234

_, axes = plt.subplots(nrows=2, ncols=1, figsize=(BASE, BASE / 1.2), sharex=True)


def set_seeds(seed: int):
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)


def store_json(d: dict, *, file: str) -> None:
    with open(file, "w") as f:
        json.dump(d, f, indent=4)


def positivity(x: np.ndarray):
    return np.all((x >= 0).reshape(len(x), -1), axis=-1)


def constant_population(x: np.ndarray):
    return np.all(np.sum(x, axis=-1) <= 100, axis=-1)


def less_20_infected(x: np.ndarray):
    return np.all(x[:, :, 1] <= 20, axis=-1)


def time_series_plot(
    x: np.ndarray, *, features_names: list, color: str, label: str, alpha=0.5
) -> None:
    k = len(features_names)

    for i, name in enumerate(features_names):
        axes[i].plot(x[:, :, i].T, alpha=alpha, color=color, linewidth=0.8)
        axes[i].set_ylabel(name)
    plt.xlabel("time")


def l1_divergence(x: np.ndarray, y: np.ndarray) -> float:
    bounds = (min(np.min(x), np.min(y)), max(np.max(x), np.max(y)))
    bins = np.linspace(-0.5, 100.5, 102)

    def density_estimator(k):
        h = np.histogram(k, density=False, range=bounds, bins=bins)[0]
        return h / np.sum(h)

    px = density_estimator(x)
    py = density_estimator(y)

    return np.sum(np.abs(px - py)) / 2


def l1_divergences(x, y):
    x = np.round(x)
    y = np.round(y)

    l = x.shape[1]
    d_S = [l1_divergence(x[:, t, 0], y[:, t, 0]) for t in range(l)]
    d_I = [l1_divergence(x[:, t, 1], y[:, t, 1]) for t in range(l)]
    return d_S, d_I


set_seeds(SEED)

SIMULATION_FILE = "paper_utils/eSIRS_S0=80_I0=10_S25=30.npy"

if not os.path.exists(SIMULATION_FILE):
    simulate_bridging()

path = sys.argv[1]
# config_module = load_module(path=path + '/config.py')


colors = {"Guidance": "orange", "Simulated": "teal"}
x = {
    "Guidance": np.load(path + "/report/guidance.npy"),
    "Simulated": np.load("paper_utils/eSIRS_S0=80_I0=10_S25=30.npy"),
}

x["Guidance"] = x["Guidance"][np.random.permutation(len(x["Guidance"]))]
x["Simulated"] = x["Simulated"][np.random.permutation(len(x["Simulated"]))]

index = {"S": 0, "I": 1}

ALPHA = 0.6
N_SAMPLES = 10
h_line = np.mean(less_20_infected(x["Guidance"])) > 0.5
bridging = np.mean(np.round(x["Guidance"][:, 0, 0]) == 80.0) > 0.5
stats = {}

for i in ("Simulated", "Guidance"):
    samples = np.round(x[i])
    sub_sample = samples[:N_SAMPLES]
    time_series_plot(
        sub_sample, features_names=["S", "I"], color=colors[i], alpha=ALPHA, label=i
    )
    stats[f"{i} positivity"] = np.mean(positivity(samples))
    stats[f"{i} constant population: "] = np.mean(constant_population(samples))

for ax in axes:
    ax.grid(alpha=0.4, zorder=-1)


if True:
    s = np.round(x["Guidance"])
    stats[f"|S(0)-80|"] = float(np.mean(abs(s[:, 0, 0] - 80)))
    stats[f"|I(0)-10|"] = float(np.mean(abs(s[:, 0, 1] - 10)))
    stats[f"|S(25)-30|"] = float(np.mean(abs(s[:, 25, 0] - 30)))

    dS, dI = l1_divergences(x["Simulated"], x["Guidance"])
    for d in ("dS", "dI"):
        for stat in ("max", "median"):
            s = eval(f"np.{stat}({d})")
            stats[f"{stat}({d[1]})"] = float(s)


legend = axes[0].legend(("Simulated", "Guidance"))
legend.legend_handles[0].set_color(colors["Simulated"])
legend.legend_handles[1].set_color(colors["Guidance"])

store_json(stats, file=path + "/plot_stats.json")
print(stats)

plt.tight_layout(pad=0.1)
plt.savefig(path + "/paper_plot.svg")
plt.show()

paper_utils/eSIRS_plot_sim_compare.py
DavideScassola/score-based-constrained-generation
import os
import sys

from PIL import Image


def make_grid(image_folder, rows, cols):
    images = []
    for filename in os.listdir(image_folder):
        if filename.endswith(".png"):
            img_path = os.path.join(image_folder, filename)
            img = Image.open(img_path)
            images.append(img)

    if len(images) < rows * cols:
        print("Not enough images in the folder.")
        return

    collage_width = images[0].width * cols
    collage_height = images[0].height * rows
    collage = Image.new("RGB", (collage_width, collage_height))

    for row in range(rows):
        for col in range(cols):
            index = row * cols + col
            collage.paste(
                images[index], (col * images[0].width, row * images[0].height)
            )

    collage.save("collage.png")
    print("Collage saved as collage.png")


if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("Usage: python make_grid.py path/to/image_folder <rows> <cols>")
    else:
        image_folder = sys.argv[1]
        rows = int(sys.argv[2])
        cols = int(sys.argv[3])
        make_grid(image_folder, rows, cols)

paper_utils/make_grid.py
DavideScassola/score-based-constrained-generation
import sys

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns

path = sys.argv[1]

# Read the data from CSV files
df_no_guidance = pd.read_csv(path + "/report/no_guidance.csv")
df_guidance = pd.read_csv(path + "/report/guidance.csv")
df_rejection_sampling = pd.read_csv(path + "/report/rejection_sampling.csv")

# Set up the figure and axes for the two subplots
fig, (ax_upper, ax_lower) = plt.subplots(
    nrows=2, ncols=1, figsize=(5.0, 2), sharex=True
)

# Set the number of bins for both subplots
bins = 50

# Plot histogram with KDE for 'no_guidance.csv' in the upper subplot
sns.histplot(
    data=df_no_guidance,
    x=df_no_guidance.columns[0],
    kde=True,
    stat="density",
    bins=bins,
    ax=ax_upper,
)
# ax_upper.set_title("Histogram of 'no_guidance.csv'")
ax_upper.legend(labels=["No guidance"])

# Calculate histogram bins for the lower plot using all data
# _, bins, _ = ax_upper.hist(df_guidance.iloc[:, 0], bins=bins, density=True)
sns.histplot(
    data=df_guidance,
    x=df_guidance.columns[0],
    kde=True,
    stat="density",
    bins=bins,
    ax=ax_lower,
    color="blue",
    alpha=0.35,
)

_, bins, _ = ax_lower.hist(df_guidance.iloc[:, 0], bins=bins, density=True, alpha=0.0)


sns.histplot(
    data=df_rejection_sampling,
    x=df_rejection_sampling.columns[0],
    kde=True,
    stat="density",
    bins=bins,
    ax=ax_lower,
    color="red",
    alpha=0.35,
)

# ax_lower.set_title("Comparison of 'guidance.csv' and 'rejection_sampling.csv'")
ax_lower.legend(labels=["Guidance", "Rejection sampling"])

# Set common labels
# fig.text(0.5, 0.04, "Values", ha="center")
# fig.text(0.04, 0.5, "Probability Density", va="center", rotation="vertical")
ax_upper.set_xlabel("")
ax_lower.set_xlabel("")

# Set x-limits of the lower subplot to match the upper subplot
lower_xlim = ax_upper.get_xlim()
ax_lower.set_xlim(lower_xlim)
# ax_upper.set_xticks([])  # Remove x-ticks from the upper subplot

plt.tight_layout(pad=0.2)
plt.savefig(path + "/paper_plot.svg")
plt.show()

paper_utils/toy_data_plot.py
DavideScassola/score-based-constrained-generation
import sys

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns

path = sys.argv[1]

# Read the data from CSV files
df_no_guidance = pd.read_csv(path + "/report/no_guidance.csv")
df_guidance = pd.read_csv(path + "/report/guidance.csv")
df_rejection_sampling = pd.read_csv(path + "/report/rejection_sampling.csv")

# select interesting columns
features = ["alcohol", "fixed acidity", "residual sugar", "citric acid"]
n_features = len(features)

# Set up the figure and axes for the two subplots
cols = 2
fig, axes = plt.subplots(
    nrows=n_features // cols, ncols=cols, figsize=(4.5, 3.4), sharex=False, sharey=False
)
axes = axes.flatten()


# Set the number of bins for both subplots
bins = int(len(df_guidance) ** 0.5)

# colors
colors = ["gray", "blue", "red"]

ALPHA = 0.35
BW_ADJUST = {
    "alcohol": 0.5,
    "fixed acidity": 0.15,
    "residual sugar": 0.5,
    "citric acid": 0.4,
}

for i in range(n_features):
    ax = axes[i]
    feature = features[i]
    bw_adjust = BW_ADJUST[feature]

    sns.kdeplot(
        data=df_no_guidance,
        x=df_no_guidance[feature],
        # kde=True,
        # stat="density",
        # bins=bins,
        ax=ax,
        color=colors[0],
        bw_adjust=0.8,
        alpha=0.5,
    )

    # _, copied_bins, _ = ax.hist(
    #    df_guidance[feature], bins=bins, density=True, alpha=0.0
    # )

    if feature == "fixed acidity":
        bw_adjust = 0.2

    sns.kdeplot(
        data=df_rejection_sampling,
        x=df_rejection_sampling[feature],
        # kde=True,
        # stat="density",
        # bins=copied_bins,
        ax=ax,
        color=colors[2],
        alpha=ALPHA,
        bw_adjust=bw_adjust,
        fill=True,
    )

    sns.kdeplot(
        data=df_guidance,
        x=df_guidance[feature],
        # kde=True,
        # stat="density",
        # bins=bins,
        ax=ax,
        color=colors[1],
        alpha=ALPHA,
        bw_adjust=bw_adjust,
        fill=True,
    )

    # ax.set_ylabel(feature)
    # ax.set_xlabel(None)
    ax.tick_params(axis="y", which="both", labelleft=False)
for i in (1, 3):
    axes[i].set_ylabel(None)

axes[-2].legend(labels=["No guidance", "Guidance", "Rej. samp."])


plt.tight_layout(pad=0.2)
plt.savefig(path + "/paper_plot.svg")
plt.show()

paper_utils/wine_plot.py
DavideScassola/score-based-constrained-generation
import sys

file = sys.argv[1]

with open(file, "r") as f:
    header = f.readline()

    for i, h in enumerate(header.split(",")):
        print(f"{i}: {h}")

scripts/show_df_info.py
DavideScassola/score-based-constrained-generation
from pathlib import Path

EXPERIMENTS_FOLDER = Path("artifacts")
MODELS_FOLDER_NAME = Path("models")
CONFIG_FILE_NAME = Path("config.py")
CONSTRAINED_GENERATIONS_FOLDER = Path("constrained_generation")
DATASET_SCRIPTS_FOLDER = Path("datasets_scripts")
MINIMUM_SAMPLING_T = 1e-4  # TODO: check this

MODELS_FOLDER = EXPERIMENTS_FOLDER / MODELS_FOLDER_NAME
LANGEVIN_CLEANING_PATIENCE = 100  # TODO: make configurable
ALLOW_LANGEVIN_CLEANING = False  # TODO: make configurable

src/constants.py
DavideScassola/score-based-constrained-generation
from dataclasses import dataclass

from src.constraints.constraint import Constraint
from src.util import load_module


@dataclass
class ConstrainedGenerationConfig:
    constraint: Constraint
    model_path: str
    generation_options: dict
    name: str | None = None
    seed: int | None = None


def getConfig(path: str) -> ConstrainedGenerationConfig:
    return load_module(str(path)).CONFIG

src/constrained_generation_config.py
DavideScassola/score-based-constrained-generation
import os
import shutil
from pathlib import Path

import src.train_config as train_config
from src.constants import (CONFIG_FILE_NAME, CONSTRAINED_GENERATIONS_FOLDER,
                           EXPERIMENTS_FOLDER)
from src.constrained_generation_config import (ConstrainedGenerationConfig,
                                               getConfig)
from src.constraints.constraint_application import apply_constraint_
from src.models.load_model import Model, load_model
from src.models.score_based.score_based_sde import ScoreBasedSde
from src.util import create_experiment_folder, set_seeds


def create_constrained_generations_folder() -> Path:
    models_folder = EXPERIMENTS_FOLDER / CONSTRAINED_GENERATIONS_FOLDER
    os.makedirs(models_folder, exist_ok=True)
    return models_folder


def store_results(
    *,
    config_file: str | Path,
    constrained_generation_config: ConstrainedGenerationConfig,
    model: Model,
):
    training_config = train_config.getConfig(
        f"{constrained_generation_config.model_path}/{CONFIG_FILE_NAME}"
    )
    training_config.generation_options = (
        constrained_generation_config.generation_options
    )
    folder: Path = Path(
        create_experiment_folder(
            path=create_constrained_generations_folder(),
            postfix=getConfig(str(config_file)).name,
        )
    )
    model.generate_report(
        path=folder,
        dataset=training_config.dataset,
        generation_options=training_config.generation_options,
        constraint=constrained_generation_config.constraint,
    )
    shutil.copyfile(config_file, folder / CONFIG_FILE_NAME)


def run_constrained_generation(config_file: str):
    constrained_generation_config = getConfig(config_file)
    if constrained_generation_config.seed:
        set_seeds(constrained_generation_config.seed)
    model = load_model(constrained_generation_config.model_path)

    if not isinstance(model, ScoreBasedSde):
        raise ValueError(
            "Currently constrained generation is allowed only for ScoreBasedSde models"
        )

    apply_constraint_(model=model, constraint=constrained_generation_config.constraint)
    store_results(
        config_file=config_file,
        constrained_generation_config=constrained_generation_config,
        model=model,
    )

src/constrained_generation_controller.py
DavideScassola/score-based-constrained-generation
from dataclasses import dataclass
from typing import Any, Callable

import torch

from src.constraints.gradient_mixer import GradientMixer, linear_interpolation


@dataclass
class Constraint:
    f: Callable[[torch.Tensor], torch.Tensor]
    gradient_mixer: GradientMixer = linear_interpolation
    strength: float = 1.0

    def __call__(self, x: torch.Tensor) -> torch.Tensor:
        return self.strength * self.f(x)

src/constraints/constraint.py
DavideScassola/score-based-constrained-generation
from typing import Callable, List

import torch

from src.constraints.constraint import Constraint
from src.models.score_based.score_based_sde import ScoreBasedSde
from src.models.score_based.score_function import ScoreFunction
from src.models.score_based.sdes.sde import VE, Sde
from src.preprocessors.preprocessor import (TensorPreprocessor,
                                            composed_inverse_transform)
from src.util import gradient

GRADIENT_CLIPPING = None  # TODO: is clipping OK?
CONSTRAINT_PREDICTION = False  # TODO: experimental


def prediction_based_constrained_score(
    *, sde: Sde, t: float, X: torch.Tensor, score_function, constraint: Constraint
) -> torch.Tensor:
    # EXPERIMENTAL!

    t_tensor = torch.full(
        size=(len(X), 1),
        fill_value=t,
        dtype=torch.float32,
        device=X.device,
    )

    X.requires_grad = True
    X.grad = torch.zeros_like(X)

    score = score_function(
        X=X, t=t_tensor
    ).detach()  # TODO: maybe use the approximated constrained score instead?

    predicted_X0 = sde.denoising(t=t_tensor, X=X, score=score)

    predicted_constraint_value = constraint(predicted_X0)
    predicted_constraint_value.sum().backward()

    # TODO: all gradients are the same!
    constrained_score = constraint.gradient_mixer(
        constraint_gradient=X.grad, score=score.detach(), t=t, sde=sde
    )

    X.grad = None
    X.requires_grad = False

    return constrained_score


def theoretical_score_limit(*, X: torch.Tensor, sde: Sde, t: float):
    """
    Let's consider inpainting for one dimension x, where the value h is imposed.
    x_t ~ N(mu(h, t), sigma(h, t)) where mu() and sigma() depend on the sde.
    Then the score of x_t is s_inp(x_t) = ( mu(h, t) - x_t ) / sigma(h, t)**2.
    |s_inp(x_t)| can be considered the limit in the value of the score, since inpainting
    is the strongest constraint. h can be an extreme value but in the gaussian range (for example +3 or -3)
    """
    # EXPERIMENTAL!
    EXTREME = 3.0
    t_tensor = torch.full((len(X),), t)
    mean1, std1 = sde.transition_kernel_mean_std(
        X=torch.full_like(X, EXTREME), t=t_tensor
    )
    mean2, std2 = sde.transition_kernel_mean_std(
        X=torch.full_like(X, -EXTREME), t=t_tensor
    )
    score1 = (mean1 - X) / std1**2
    score2 = (mean2 - X) / std2**2
    limit = torch.maximum(score1.abs(), score2.abs())
    return limit


def get_constrained_score(
    score_function: ScoreFunction, constraint: Constraint
) -> ScoreFunction:
    class ConstrainedScoreFunction(ScoreFunction):
        def __init__(
            self, *, original_score_function: ScoreFunction, constraint: Constraint
        ) -> None:
            self.original_score_function = original_score_function
            super().__init__(
                sde=self.original_score_function.sde,
                shape=self.original_score_function.shape,
                device=self.original_score_function.device,
            )
            self.constraint = constraint

            # TODO: check if appropriate
            for param in self.parameters():
                param.requires_grad = False

            self.train(False)

        def forward(self, *, X: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
            return self.original_score_function.forward(X=X, t=t)

        def build_model(self, **hyperparameters) -> torch.nn.Module:
            return self.original_score_function.model

        def noised_constraint_score_estimate(
            self, *, X: torch.Tensor, t: float, n: int
        ):
            mean, std = self.sde.transition_kernel_mean_std(
                X=X,
                t=torch.full(
                    size=(len(X), 1), fill_value=t, dtype=torch.float32, device=X.device
                ),
            )
            shape = [n] + list(X.shape)
            samples = torch.rand(size=shape, device=X.device) * std + mean
            # TODO: check reshaping
            return constraint.strength * gradient(
                f=constraint.f, X=samples.reshape([-1] + shape[2:])
            ).reshape(shape).mean(dim=0)

        def __call__(self, *, X: torch.Tensor, t: float) -> torch.Tensor:
            """
            p_constrained(x, t) = p(x, t)*exp(lambda*f(x, t))

            score(p_constrained(x, t)) =
            grad{log[p(x, t)*exp(lambda*f(x, t))]} =
            grad{logp(x, t) + lambda*f(x, t)]} =
            grad{logp(x, t) + lambda*f(x, t)]} =
            grad{logp(x, t)} + lambda*grad{f(x, t)} =
            s(x, t) + lambda*grad{f(x, t)}
            """

            # TODO: osserva come score dei dati cambia col tempo
            # TODO: applica riscalo dello score dei dati al gradiente del vincolo
            if self.constraint is None:
                return super().__call__(X=X, t=t)

            if CONSTRAINT_PREDICTION:
                return prediction_based_constrained_score(
                    sde=self.sde,
                    t=t,
                    X=X,
                    score_function=super().__call__,
                    constraint=self.constraint,
                )

            else:
                constraint_gradient = (
                    gradient(f=self.constraint.f, X=X) * self.constraint.strength
                )
                if GRADIENT_CLIPPING:
                    constraint_gradient = torch.clip(
                        gradient(f=self.constraint.f, X=X) * self.constraint.strength,
                        min=-GRADIENT_CLIPPING,
                        max=GRADIENT_CLIPPING,
                    )

                return self.constraint.gradient_mixer(
                    score=super().__call__(X=X, t=t),
                    constraint_gradient=constraint_gradient,
                    t=t,
                    sde=self.sde,
                    X=X,
                    constraint=self.constraint,
                )

    return ConstrainedScoreFunction(
        original_score_function=score_function, constraint=constraint
    )


def include_inverse_preprocessing(
    *, constraint: Constraint, preprocessors: List[TensorPreprocessor]
) -> Constraint:
    return Constraint(
        f=lambda x: constraint.f(
            composed_inverse_transform(x, preprocessors=preprocessors)
        ),
        strength=constraint.strength,
        gradient_mixer=constraint.gradient_mixer,
    )


def apply_constraint_(*, model: ScoreBasedSde, constraint: Constraint) -> ScoreBasedSde:
    model.score_model = get_constrained_score(
        score_function=model.score_model,
        constraint=include_inverse_preprocessing(
            constraint=constraint, preprocessors=model.tensor_preprocessors
        ),
    )
    return model

src/constraints/constraint_application.py
DavideScassola/score-based-constrained-generation
from abc import ABC, abstractmethod

import numpy as np
import torch

from src.constants import MINIMUM_SAMPLING_T
from src.models.score_based.sdes.sde import Sde
from src.util import gradient


def fix_bounds(g):
    """
    Rescales a function g(t) such that g(0)=1 and g(1)=0
    """
    if g(0.0) == 1 and g(1.0) == 0.0:
        return g
    h = -g(1.0) * (g(0.0) - g(1.0))
    m = 1.0 / (g(0.0) - g(1.0))
    return lambda t: g(t) * m + h


class GradientMixer(ABC):
    @abstractmethod
    def __call__(
        self,
        *,
        constraint_gradient: torch.Tensor,
        score: torch.Tensor,
        t: float,
        sde: Sde,
        **kwargs
    ) -> torch.Tensor:
        pass


class NoConstraint(GradientMixer):
    def __call__(
        self,
        *,
        constraint_gradient: torch.Tensor,
        score: torch.Tensor,
        t: float,
        sde: Sde,
        **kwargs
    ) -> torch.Tensor:
        return score


class TimeWeighting(GradientMixer):
    def __init__(self, g, rescale: bool = True) -> None:
        self.g = fix_bounds(g) if rescale else g

    def __call__(
        self,
        *,
        constraint_gradient: torch.Tensor,
        score: torch.Tensor,
        t: float,
        sde: Sde,
        **kwargs
    ) -> torch.Tensor:
        return constraint_gradient * self.g(t) + score


class LogisticTimeWeighting(TimeWeighting):
    def __init__(self, *, k: float = 10.0, x0: float = 0.5) -> None:
        g = lambda t: 1 - 1 / (1 + np.exp(-(t - x0) * k))
        super().__init__(g, rescale=True)


def expected_score(*, n: int, t: float, sde: Sde, X: torch.Tensor, constraint):
    mean, std = sde.transition_kernel_mean_std(
        X=X,
        t=torch.full(
            size=(len(X), 1), fill_value=t, dtype=torch.float32, device=X.device
        ),
    )
    shape = [n] + list(X.shape)
    samples = torch.rand(size=shape, device=X.device) * std + mean
    # TODO: check reshaping
    return constraint.strength * gradient(
        f=constraint.f, X=samples.reshape([-1] + shape[2:])
    ).reshape(shape).mean(dim=0)


class NoisedConstraint(GradientMixer):
    def __init__(self, n: int) -> None:
        self.n = n

    def __call__(
        self,
        *,
        constraint_gradient: torch.Tensor,
        score: torch.Tensor,
        t: float,
        sde: Sde,
        X: torch.Tensor,
        constraint,
        **kwargs
    ) -> torch.Tensor:
        return (
            expected_score(constraint=constraint, t=t, n=self.n, X=X, sde=sde) + score
        )


class GradientClipping(GradientMixer):
    def __init__(self, k: float) -> None:
        super().__init__()
        self.k = abs(k)

    def __call__(
        self,
        *,
        constraint_gradient: torch.Tensor,
        score: torch.Tensor,
        t: float,
        sde: Sde,
        **kwargs
    ) -> torch.Tensor:
        _, std = sde.transition_kernel_mean_std(X=torch.tensor(0.0), t=torch.tensor(t))
        cap = (self.k / (std**2)).to(device=constraint_gradient.device)
        clipped_constraint_grad = torch.clip(constraint_gradient, min=-cap, max=cap)
        return score + clipped_constraint_grad


class NoiseWeighting(GradientMixer):
    def __init__(self, min_t: float = MINIMUM_SAMPLING_T) -> None:
        super().__init__()
        self.min_t = min_t
        self.min_sigma = None

    def __call__(
        self,
        *,
        constraint_gradient: torch.Tensor,
        score: torch.Tensor,
        t: float,
        sde: Sde,
        **kwargs
    ) -> torch.Tensor:
        if not self.min_sigma:
            self.min_sigma = sde.transition_kernel_mean_std(
                X=torch.tensor(0.0), t=torch.tensor(self.min_t)
            )[1]

        sigma = sde.transition_kernel_mean_std(X=torch.tensor(0.0), t=torch.tensor(t))[
            1
        ]
        return score + constraint_gradient * self.min_sigma / sigma


class ExpNoiseWeighting(GradientMixer):
    def __init__(self, min_t: float = MINIMUM_SAMPLING_T) -> None:
        super().__init__()
        self.min_sigma = None

    def __call__(
        self,
        *,
        constraint_gradient: torch.Tensor,
        score: torch.Tensor,
        t: float,
        sde: Sde,
        **kwargs
    ) -> torch.Tensor:
        if not self.min_sigma:
            self.min_sigma = sde.transition_kernel_mean_std(
                X=torch.tensor(0.0), t=torch.tensor(0)
            )[1]

        sigma = sde.transition_kernel_mean_std(X=torch.tensor(0.0), t=torch.tensor(t))[
            1
        ]
        return score + constraint_gradient * torch.exp(-sigma + self.min_sigma)


class SNRWeighting(GradientMixer):
    def __call__(
        self,
        *,
        constraint_gradient: torch.Tensor,
        score: torch.Tensor,
        t: float,
        sde: Sde,
        **kwargs
    ) -> torch.Tensor:
        return score + constraint_gradient * sde.snr(t)


class SNRC(GradientMixer):
    def __init__(self, k: float) -> None:
        super().__init__()
        self.k = abs(k)

    def __call__(
        self,
        *,
        constraint_gradient: torch.Tensor,
        score: torch.Tensor,
        t: float,
        sde: Sde,
        **kwargs
    ) -> torch.Tensor:
        _, std = sde.transition_kernel_mean_std(X=torch.tensor(0.0), t=torch.tensor(t))
        cap = (self.k / (std**2)).to(device=constraint_gradient.device)
        weighted_constraint = constraint_gradient * sde.snr(t)
        clipped_constraint_grad = torch.clip(weighted_constraint, min=-cap, max=cap)
        return score + clipped_constraint_grad


linear_interpolation = TimeWeighting(g=lambda t: 1 - t)
exponential_interpolation = TimeWeighting(g=lambda t: np.exp(-5.0 * t))
logistic_interpolation = LogisticTimeWeighting(k=10.0)
constant = TimeWeighting(g=lambda t: 1.0, rescale=False)
noise_weighting = NoiseWeighting()
exp_noise_weighting = ExpNoiseWeighting()
snr = SNRWeighting()

src/constraints/gradient_mixer.py
DavideScassola/score-based-constrained-generation
import math

import torch

"""
common numerical trick: I use two different but analitically equivalent functions in order to avoid computing large exponentials in different regions
"""

MAXGRAD = 50
MAX_NEG = 20


def numerically_stable_log_logistic(x: torch.Tensor) -> torch.Tensor:
    return torch.where(
        x > 0, -torch.log(1 + torch.exp(-x)), -torch.log(1 + torch.exp(x)) + x
    )


def numerically_stable_log_logistic_derivative(x: torch.Tensor) -> torch.Tensor:
    return torch.where(
        x > 0, torch.exp(-x) / (torch.exp(-x) + 1), 1 / (torch.exp(x) + 1)
    )


class LogLogistic(torch.autograd.Function):
    """
    Numerically stable implementation of the log of the logistic function ( -log(1 + exp(-x)) )
    """

    @staticmethod
    def forward(ctx, input: torch.Tensor) -> torch.Tensor:
        ctx.save_for_backward(input)
        return numerically_stable_log_logistic(input)

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
        (input,) = ctx.saved_tensors
        return grad_output * numerically_stable_log_logistic_derivative(input)


log_logistic = LogLogistic.apply

src/constraints/numerically_stable_functions.py
DavideScassola/score-based-constrained-generation
from abc import ABC, abstractclassmethod

import torch

LOG_EPS = 1e-6
STABLE_PRODUCT_EPS = 1e-2
log_logistic = torch.nn.LogSigmoid()


class RealLogic(ABC):
    @staticmethod
    def greater(
        x1: torch.Tensor, x2: torch.Tensor | float, grad: float = 1.0
    ) -> torch.Tensor:
        return torch.where(x1 > x2, 0.0, x1 - x2) * grad

    @staticmethod
    def smaller(
        x1: torch.Tensor, x2: torch.Tensor | float, grad: float = 1.0
    ) -> torch.Tensor:
        return torch.where(x1 < x2, 0.0, x2 - x1) * grad

    @staticmethod
    def equal(
        x1: torch.Tensor, x2: torch.Tensor | float, k: float = 1.0
    ) -> torch.Tensor:
        return -k * torch.square(x1 - x2)

    @classmethod
    def in_(
        cls,
        x: torch.Tensor,
        inf: torch.Tensor | float,
        sup: torch.Tensor | float,
        grad: float = 1.0,
    ) -> torch.Tensor:
        return cls.and_(cls.greater(x, inf, grad), cls.smaller(x, sup, grad))  # type: ignore

    @abstractclassmethod
    def and_(cls, *x: torch.Tensor) -> torch.Tensor:  # type: ignore
        pass

    @abstractclassmethod
    def or_(cls, *x: torch.Tensor) -> torch.Tensor:  # type: ignore
        pass

    @abstractclassmethod
    def implies(cls, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        return cls.or_(cls.not_(x1), x2)

    @classmethod
    def not_(cls, x: torch.Tensor) -> torch.Tensor:  # type: ignore
        pass


class Product(RealLogic):
    """
    Inspired from Badreddine, et al. (2020). Logic Tensor Networks, https://arxiv.org/abs/2012.13635

    For the conjunction and negation
    """

    DEFAULT_GRAD = 50.0

    @staticmethod
    def greater(
        x1: torch.Tensor, x2: torch.Tensor | float, grad: float = DEFAULT_GRAD
    ) -> torch.Tensor:
        return log_logistic(grad * (x1 - x2))  # type: ignore

    @staticmethod
    def smaller(
        x1: torch.Tensor, x2: torch.Tensor | float, grad: float = DEFAULT_GRAD
    ) -> torch.Tensor:
        return log_logistic(grad * (x2 - x1))  # type: ignore

    @classmethod
    def and_(cls, *x: torch.Tensor) -> torch.Tensor:
        """
        If multiple tensors are provided, the element wise 'and' is returned
        If a single tensor is given, then the logical 'and' reduces all dimensions except the first
        """
        if len(x) == 1:
            return x[0].flatten(1).sum(1)
        return torch.stack(x).sum(0)

    @classmethod
    def or_(cls, *x: torch.Tensor, eps=LOG_EPS) -> torch.Tensor:
        if len(x) == 1:
            return x[0]
        if len(x) == 2:
            return torch.log(
                torch.exp(x[0]) + torch.exp(x[1]) - torch.exp(x[0] + x[1]) + eps
            )
        return cls.or_(cls.or_(x[0], x[1]), *x[2:])

    @classmethod
    def not_(cls, x: torch.Tensor) -> torch.Tensor:
        """
        This doesn't work numerically, very hard problem to solve
        """
        return torch.log(1 - torch.exp(x) + LOG_EPS)

    @classmethod
    def in_(
        cls,
        x: torch.Tensor,
        inf: torch.Tensor | float,
        sup: torch.Tensor | float,
        grad: float = DEFAULT_GRAD,
    ) -> torch.Tensor:
        return cls.and_(cls.greater(x, inf, grad), cls.smaller(x, sup, grad))  # type: ignore

    @classmethod
    def equal(
        cls, x1: torch.Tensor, x2: torch.Tensor | float, grad: float = DEFAULT_GRAD
    ) -> torch.Tensor:
        return cls.in_(x1, inf=x2, sup=x2, grad=grad)


class Minimum(RealLogic):
    """
    Inspired from Badreddine, et al. (2020). Logic Tensor Networks, https://arxiv.org/abs/2012.13635
    """

    @classmethod
    def and_(cls, *x: torch.Tensor) -> torch.Tensor:
        return torch.stack(x).amin(0)

    @classmethod
    def or_(cls, *x: torch.Tensor) -> torch.Tensor:
        return torch.stack(x).amax(0)


class Classic(RealLogic):
    """
    Safe logic that I used at the beginning (for aggreating the logical 'and' I used the sum instead, in order to avoid 'single-passing')
    """

    @classmethod
    def and_(cls, *x: torch.Tensor) -> torch.Tensor:
        """
        If multiple tensors are provided, the element wise 'and' is returned
        If a single tensor is given, then the logical 'and' reduces all dimensions except the first
        """
        if len(x) == 1:
            return x[0].flatten(1).sum(1)
        return torch.stack(x).sum(0)

    @classmethod
    def or_(cls, *x: torch.Tensor) -> torch.Tensor:
        return torch.stack(x).amax(0)

    @classmethod
    def implies(cls, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        return torch.where(x1 >= 0.0, x2, 0.0)

    @classmethod
    def not_(cls, x: torch.Tensor) -> torch.Tensor:
        raise NotImplementedError("Not implemented, use negation only on atoms")


class StableProduct(Product):
    """
    Inspired from Badreddine, et al. (2020). Logic Tensor Networks, https://arxiv.org/abs/2012.13635

    Product logic with numerical corrections in order to avoid gradient issues, but in this case this could be unnecessary since we consider gradient of the log.
    """

    @staticmethod
    def pi0(x: torch.Tensor) -> torch.Tensor:
        """This transformation avoids a value of 0"""
        return torch.log((1 - STABLE_PRODUCT_EPS) * torch.exp(x) + STABLE_PRODUCT_EPS)

    @staticmethod
    def pi1(x: torch.Tensor) -> torch.Tensor:
        """This transformation avoids a value of 1"""
        return torch.log((1 - STABLE_PRODUCT_EPS) * torch.exp(x))

    @classmethod
    def and_(cls, *x: torch.Tensor) -> torch.Tensor:
        return super().and_(*(StableProduct.pi0(xi) for xi in x))

    @classmethod
    def or_(cls, *x: torch.Tensor) -> torch.Tensor:
        return super().or_(*(StableProduct.pi1(xi) for xi in x))

    @classmethod
    def implies(cls, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        return cls.implies(StableProduct.pi0(x1), StableProduct.pi1(x2))

src/constraints/real_logic.py
DavideScassola/score-based-constrained-generation
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from pandas import DataFrame
from torch import Tensor
from torchvision import datasets, transforms
from torchvision.transforms import ToTensor

from src.constants import DATASET_SCRIPTS_FOLDER
from src.util import execute_python_script


def get_subset_slice(dataset, train_prop: float, train: bool = True) -> slice:
    cut = int(len(dataset) * train_prop)
    return slice(None, cut) if train else slice(cut, None)


def generate_from_script(dataset_name: str):
    script = DATASET_SCRIPTS_FOLDER / Path(dataset_name + ".py")
    if not script.is_file():
        raise ValueError(f"Couldn't find script for generating {dataset_name} dataset")
    print(f"Generating dataset from {script} script ...")
    execute_python_script(script)


def generate_if_necessary(path: str | Path) -> None:
    path = Path(path)
    if not path.is_file():
        print(f"dataset {path} not found")
        generate_from_script(path.stem)


def load_csv(path: str | Path) -> pd.DataFrame:
    return pd.read_csv(path)


def load_npy(path: str | Path) -> Tensor:
    return torch.from_numpy(np.load(path)).float()


def loader2tensor(d) -> torch.Tensor:
    # TODO: slow and bad solution, one shouldn't convert back to a single tensor
    return torch.concat([s[0] for s in d])


def get_MNIST(*, train: bool):
    # TODO: transformation is not OK for every model
    print("loading MNIST...")
    return loader2tensor(
        datasets.MNIST(
            root="data",
            train=train,
            transform=transforms.Compose(
                [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
            ),
            download=True,
        )
    )


DEFAULT_DATASETS = {"MNIST": get_MNIST}


@dataclass
class Dataset:
    path: str | Path
    train_proportion: float = 1.0

    def get(self, train=True) -> DataFrame | Tensor:
        if self.path in DEFAULT_DATASETS:
            return DEFAULT_DATASETS[self.path](train=train)

        generate_if_necessary(self.path)
        ext = Path(self.path).suffix[1:]
        data = eval(f"load_{ext}")(self.path)
        s = get_subset_slice(data, self.train_proportion, train=train)
        if isinstance(data, DataFrame):
            return data.iloc[s]
        else:
            return data[s]

src/data.py
DavideScassola/score-based-constrained-generation
import os
from pathlib import Path
from typing import List

import torch

from src.data import Dataset
from src.preprocessors.preprocessor import (TensorPreprocessor,
                                            composed_fit_transform,
                                            composed_inverse_transform)
from src.report import *

from .model import Model


class ImageModel(Model):
    def __init__(self, *, tensor_preprocessors: List[TensorPreprocessor] = []) -> None:
        self.tensor_preprocessors = tensor_preprocessors
        self.preprocessors = self.tensor_preprocessors

    def train(self, dataset: Dataset):
        X = dataset.get(train=True)
        if not isinstance(X, torch.Tensor):
            raise ValueError(f"dataset {dataset} should be a tensor")
        X = composed_fit_transform(X, preprocessors=self.tensor_preprocessors)
        self._train(X)

    def generate(self, n_samples: int, **kwargs) -> torch.Tensor:
        samples = self._generate(n_samples, **kwargs)
        return composed_inverse_transform(
            samples, preprocessors=self.tensor_preprocessors
        )

    def generate_report(
        self, *, path: str | Path, dataset: Dataset, generation_options: dict
    ):
        report_folder = path / Path(REPORT_FOLDER_NAME)
        os.makedirs(report_folder, exist_ok=False)

        samples = self.generate(**generation_options).cpu()
        images_folder = str(path) + "/samples"
        os.makedirs(images_folder)

        store_images(samples, folder=images_folder)

        plt.imshow(samples[0])
        plt.colorbar()
        plt.savefig(f"image.{IMAGE_FORMAT}")
        plt.close()

src/models/image_model.py
DavideScassola/score-based-constrained-generation
from pathlib import Path

from src.constants import CONFIG_FILE_NAME
from src.models.model import Model
from src.train_config import getConfig


def load_model(model_path: str) -> Model:
    train_config = getConfig(Path(model_path) / CONFIG_FILE_NAME)
    train_config.model.load_(model_path)
    return train_config.model

src/models/load_model.py
DavideScassola/score-based-constrained-generation
import os
from abc import ABC, abstractmethod
from pathlib import Path
from typing import List

import pandas as pd
import torch

from src.data import Dataset
from src.preprocessors.preprocessor import Preprocessor


class Model(ABC):
    MODEL_FOLDER_NAME = Path("model")
    PARAMS_FILE = Path("params.json")

    def __init__(self, *, preprocessors: List[Preprocessor] = []) -> None:
        super().__init__()
        self.preprocessors = preprocessors

    def store(self, experiment_path: str | Path):
        model_path = self.get_model_folder(experiment_path)
        os.makedirs(model_path)
        self.store_preprocessors(model_path)
        self._store(model_path)

    def load_(self, experiment_path: str | Path):
        model_path = self.get_model_folder(experiment_path)
        self.load_preprocessors(model_path)
        self._load_(model_path)

    def get_model_folder(self, experiment_path: str | Path) -> Path:
        return Path(experiment_path) / self.MODEL_FOLDER_NAME

    def params_file(self, model_path: str | Path):
        return Path(model_path) / self.PARAMS_FILE

    def store_preprocessors(self, model_path: str | Path):
        for p in self.preprocessors:
            p.store(model_path)

    def load_preprocessors(self, model_path: str | Path):
        for p in self.preprocessors:
            p.load_(model_path)

    def specific_report_plots(self, path: Path):
        pass

    @abstractmethod
    def train(self, dataset: Dataset):
        pass

    @abstractmethod
    def generate(self, n_samples: int, **kwargs) -> pd.DataFrame:
        pass

    @abstractmethod
    def _train(self, X: torch.Tensor):
        pass

    @abstractmethod
    def _generate(self, n_samples: int) -> torch.Tensor:
        pass

    @abstractmethod
    def _store(self, model_path: str | Path):
        pass

    @abstractmethod
    def _load_(self, model_path: str | Path):
        pass

    @abstractmethod
    def generate_report(
        self, *, path: str | Path, dataset: Dataset, generation_options: dict, **kwargs
    ):
        pass

src/models/model.py
DavideScassola/score-based-constrained-generation
import numpy as np
import torch

from .model import Model


class MultivariateGaussian(Model):
    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)

    def fit_multivariate_gaussian(self, X: torch.Tensor):
        X = X.numpy()
        self.mean = np.mean(X, axis=0)
        self.cov = np.cov(X.T, ddof=1)
        if len(self.cov.shape) == 0:
            self.cov = self.cov.reshape(1, 1)

    def sample_multivariate_gaussian(self, n: int):
        return np.random.multivariate_normal(self.mean, self.cov, n)

    def _train(self, X: torch.Tensor):
        self.fit_multivariate_gaussian(X)

    # TODO: change to tensor?
    def _generate(self, n_samples: int = 1000) -> torch.Tensor:
        return torch.from_numpy(self.sample_multivariate_gaussian(n_samples))

    def store(self, path: str):
        # TODO
        pass

    def load_(self, path: str):
        # TODO
        pass

    def _store(self, path: str):
        # TODO
        pass

    def _load_(self, path: str):
        # TODO
        pass

src/models/multivariate_gaussian.py
DavideScassola/score-based-constrained-generation
import os

from src.data import Dataset
from src.models.image_model import ImageModel
from src.report import *

from .score_based_sde import ScoreBasedSde


class ImageScoreBasedSde(ScoreBasedSde, ImageModel):
    def generate_report(
        self,
        *,
        path: str | Path,
        dataset: Dataset,
        generation_options: dict,
        constraint: Constraint | None = None,
    ):
        report_folder = path / Path(REPORT_FOLDER_NAME)
        os.makedirs(report_folder, exist_ok=False)

        samples = self.generate(**generation_options).cpu()
        images_folder = str(path) + "/samples"
        os.makedirs(images_folder)

        store_images(samples, folder=images_folder)

        # if constraint:
        #    print(np.exp(constraint.f(samples).numpy()))
        plt.imshow(samples[0])
        plt.colorbar()
        plt.savefig(f"image.{IMAGE_FORMAT}")
        plt.close()

src/models/score_based/image_score_based_sde.py
DavideScassola/score-based-constrained-generation
import torch
from torch import nn

from src.models.score_based.nn.positional_embedding import PositionalEmbedding


class Block(nn.Module):
    """Residual block for the denoising model."""

    def __init__(self, size: int):
        super().__init__()

        self.layer = nn.Linear(size, size)
        self.activation = nn.GELU()

    def forward(self, x: torch.Tensor):
        return x + self.activation(self.layer(x))


class MLP(nn.Module):
    def __init__(
        self,
        n_features,
        hidden_size: int = 128,
        hidden_layers: int = 3,
        emb_size: int = 128,
        time_emb: str = "sinusoidal",
        input_emb: str = "sinusoidal",
    ):
        """Denoising NN with positional embeddings.

        Args:
            hidden_size: Hidden size of blocks.
            hidden_layers: Number of hidden layers.
            emb_size: Size of positional embeddings.
            time_emb: Type of temporal embedding.
            input_emb: Type of input embedding.

        """
        super().__init__()
        self.n_features = n_features
        self.time_mlp = PositionalEmbedding(emb_size, emb_type=time_emb)
        self.input_mlps = [
            PositionalEmbedding(emb_size, emb_type=input_emb, scale=25.0)
            for _ in range(self.n_features)
        ]

        concat_size = len(self.time_mlp.layer)
        for mlp_layer in self.input_mlps:
            concat_size += len(mlp_layer.layer)
        layers = [
            nn.Linear(concat_size, hidden_size),
            nn.GELU(),
            nn.BatchNorm1d(hidden_size),
        ]
        for _ in range(hidden_layers):
            layers.append(Block(hidden_size))
            layers.append(nn.BatchNorm1d(hidden_size))
        layers.append(nn.Linear(hidden_size, self.n_features))
        self.sequential = nn.Sequential(*layers)

    def forward(self, x, t):
        t_emb = self.time_mlp(t.squeeze())
        x = torch.cat(
            [mlp(x[:, idx]) for idx, mlp in enumerate(self.input_mlps)] + [t_emb],
            dim=-1,
        )
        x = self.sequential(x)
        return x

src/models/score_based/nn/gine_mlp.py
DavideScassola/score-based-constrained-generation
from dataclasses import dataclass
from typing import Callable

import torch


@dataclass
class Optimization:
    epochs: int
    batch_size: int
    optimizer_class: Callable
    optimizer_hyperparameters: dict
    gradient_clipping: float | None = None
    parameters_momentum: float | None = None

    def build_optimizer(self, params) -> torch.optim.Optimizer:
        return self.optimizer_class(params, **self.optimizer_hyperparameters)

src/models/score_based/nn/optimization.py
DavideScassola/score-based-constrained-generation
# based on https://github.dev/yang-song/score_sde_pytorch
# Partially based on: https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/training/moving_averages.py

# from __future__ import division
# from __future__ import unicode_literals

import torch


class ParametersEMA:
    """
    Maintains (exponential) moving average of a set of parameters.
    """

    def __init__(self, parameters, decay, use_num_updates=True):
        """
        Args:
          parameters: Iterable of `torch.nn.Parameter`; usually the result of
            `model.parameters()`.
          decay: The exponential decay.
          use_num_updates: Whether to use number of updates when computing
            averages.
        """
        if decay < 0.0 or decay > 1.0:
            raise ValueError("Decay must be between 0 and 1")
        self.decay = decay
        self.num_updates = 0 if use_num_updates else None
        self.shadow_params = [p.clone().detach() for p in parameters if p.requires_grad]
        self.collected_params = []

    def update(self, parameters):
        """
        Update currently maintained parameters.

        Call this every time the parameters are updated, such as the result of
        the `optimizer.step()` call.

        Args:
          parameters: Iterable of `torch.nn.Parameter`; usually the same set of
            parameters used to initialize this object.
        """
        decay = self.decay
        if self.num_updates is not None:
            self.num_updates += 1
            decay = min(decay, (1 + self.num_updates) / (10 + self.num_updates))
        one_minus_decay = 1.0 - decay
        with torch.no_grad():
            parameters = [p for p in parameters if p.requires_grad]
            for s_param, param in zip(self.shadow_params, parameters):
                s_param.sub_(one_minus_decay * (s_param - param))

    def copy_to(self, parameters):
        """
        Copy current parameters into given collection of parameters.

        Args:
          parameters: Iterable of `torch.nn.Parameter`; the parameters to be
            updated with the stored moving averages.
        """
        parameters = [p for p in parameters if p.requires_grad]
        for s_param, param in zip(self.shadow_params, parameters):
            if param.requires_grad:
                param.data.copy_(s_param.data)

    def store(self, parameters):
        """
        Save the current parameters for restoring later.

        Args:
          parameters: Iterable of `torch.nn.Parameter`; the parameters to be
            temporarily stored.
        """
        self.collected_params = [param.clone() for param in parameters]

    def restore(self, parameters):
        """
        Restore the parameters stored with the `store` method.
        Useful to validate the model with EMA parameters without affecting the
        original optimization process. Store the parameters before the
        `copy_to` method. After validation (or model saving), use this to
        restore the former parameters.

        Args:
          parameters: Iterable of `torch.nn.Parameter`; the parameters to be
            updated with the stored parameters.
        """
        for c_param, param in zip(self.collected_params, parameters):
            param.data.copy_(c_param.data)

    def state_dict(self):
        return dict(
            decay=self.decay,
            num_updates=self.num_updates,
            shadow_params=self.shadow_params,
        )

    def load_state_dict(self, state_dict):
        self.decay = state_dict["decay"]
        self.num_updates = state_dict["num_updates"]
        self.shadow_params = state_dict["shadow_params"]

src/models/score_based/nn/parameters_ema.py
DavideScassola/score-based-constrained-generation
import torch
from torch import nn


class SinusoidalEmbedding(nn.Module):
    def __init__(self, size: int, scale: float = 1.0):
        super().__init__()
        self.size = size
        self.scale = scale

    def forward(self, x: torch.Tensor):
        x = x * self.scale
        half_size = self.size // 2
        emb = torch.log(torch.Tensor([10000.0])) / (half_size - 1)
        emb = torch.exp(-emb * torch.arange(half_size)).to(x.device)
        emb = x.unsqueeze(-1) * emb.unsqueeze(0)
        emb = torch.cat((torch.sin(emb), torch.cos(emb)), dim=-1)
        return emb

    def __len__(self):
        return self.size


class LinearEmbedding(nn.Module):
    def __init__(self, size: int, scale: float = 1.0):
        super().__init__()
        self.size = size
        self.scale = scale

    def forward(self, x: torch.Tensor):
        x = x / self.size * self.scale
        return x.unsqueeze(-1)

    def __len__(self):
        return 1


class LearnableEmbedding(nn.Module):
    def __init__(self, size: int):
        super().__init__()
        self.size = size
        self.linear = nn.Linear(1, size)

    def forward(self, x: torch.Tensor):
        return self.linear(x.unsqueeze(-1).float() / self.size)

    def __len__(self):
        return self.size


class PositionalEmbedding(nn.Module):
    """Positional embeddings"""

    def __init__(self, size: int, emb_type: str, **kwargs):
        super().__init__()

        if emb_type == "sinusoidal":
            self.layer = SinusoidalEmbedding(size, **kwargs)
        elif emb_type == "linear":
            self.layer = LinearEmbedding(size, **kwargs)
        elif emb_type == "learnable":
            self.layer = LearnableEmbedding(size)
        else:
            raise ValueError(f"Unknown positional embedding type: {emb_type}")

    def forward(self, x: torch.Tensor):
        return self.layer(x)

src/models/score_based/nn/positional_embedding.py
DavideScassola/score-based-constrained-generation
import math
from typing import Callable, List

import torch
import torch.nn.functional as F
from torch import nn


class FlexibleResidualBlock(nn.Module):
    """
    This block is a simple layer of an MLP if the input and output dimensions are different,
    otherwise it's a residual block, where the input is summed to the output
    """

    def __init__(
        self, *, input_size: int, output_size: int, activation: nn.Module
    ) -> None:
        super().__init__()
        self.activation = activation()
        self.linear = nn.Linear(in_features=input_size, out_features=output_size)
        residual = lambda x: x + self.activation(self.linear(x))
        self.f = (
            residual
            if input_size == output_size
            else lambda x: self.activation(self.linear(x))
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.f(x)


def sin_cos_embedding(t: torch.Tensor) -> torch.Tensor:
    x = t.reshape(-1, 1) * 2 * torch.pi
    return torch.cat([torch.cos(x), torch.sin(x)], dim=-1)


def classic_embedding(t, embedding_dim=20, max_positions=10000) -> torch.Tensor:
    timesteps = t.flatten() * 1000
    assert len(timesteps.shape) == 1  # and timesteps.dtype == tf.int32
    half_dim = embedding_dim // 2
    # magic number 10000 is from transformers
    emb = math.log(max_positions) / (half_dim - 1)
    emb = torch.exp(
        torch.arange(half_dim, dtype=torch.float32, device=timesteps.device) * -emb
    )
    emb = timesteps.float()[:, None] * emb[None, :]
    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
    if embedding_dim % 2 == 1:  # zero pad
        emb = F.pad(emb, (0, 1), mode="constant")
    assert emb.shape == (timesteps.shape[0], embedding_dim)
    return emb


class TimeResidualMLP(torch.nn.Module):
    def __init__(
        self,
        *,
        in_channels: int,
        out_channels: int,
        hidden_channels: List[int],
        activation_layer: nn.Module,
        time_embedding: Callable | None = classic_embedding
    ):
        super().__init__()
        # self.time_embedding = sin_cos_embedding
        # self.time_embedding = lambda x: x
        self.time_embedding = time_embedding
        t_emb_dim = (
            self.time_embedding(torch.tensor([[1]])).shape[-1]
            if self.time_embedding
            else 0
        )

        concat_size = in_channels + t_emb_dim
        layers = [
            torch.nn.BatchNorm1d(num_features=concat_size),
            FlexibleResidualBlock(
                input_size=concat_size,
                output_size=hidden_channels[0],
                activation=activation_layer,
            ),
        ]

        for i in range(len(hidden_channels) - 1):
            layers.append(
                FlexibleResidualBlock(
                    input_size=hidden_channels[i],
                    output_size=hidden_channels[i + 1],
                    activation=activation_layer,
                )
            )
        layers.append(
            torch.nn.Linear(
                in_features=hidden_channels[-1], out_features=out_channels, bias=False
            )
        )
        # TODO: bias = False since the expected value of the score is 0, but not sure
        self.sequential = torch.nn.Sequential(*layers)

    def forward(self, *, X: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        if self.time_embedding:
            t_emb = self.time_embedding(t)
            X = torch.cat((X, t_emb), dim=-1)
        return self.sequential(X)

src/models/score_based/nn/time_residual_mlp.py
DavideScassola/score-based-constrained-generation
import torch
from torch import nn


class MyBlock(nn.Module):
    def __init__(
        self,
        shape,
        in_c,
        out_c,
        kernel_size=3,
        stride=1,
        padding=1,
        activation=None,
        normalize=True,
    ):
        super(MyBlock, self).__init__()
        self.ln = nn.LayerNorm(shape)
        self.conv1 = nn.Conv2d(in_c, out_c, kernel_size, stride, padding)
        self.conv2 = nn.Conv2d(out_c, out_c, kernel_size, stride, padding)
        self.activation = nn.SiLU() if activation is None else activation
        self.normalize = normalize

    def forward(self, x):
        out = self.ln(x) if self.normalize else x
        out = self.conv1(out)
        out = self.activation(out)
        out = self.conv2(out)
        out = self.activation(out)
        return out


def sinusoidal_embedding(n, d):
    # Returns the standard positional embedding
    embedding = torch.zeros(n, d)
    wk = torch.tensor([1 / 10_000 ** (2 * j / d) for j in range(d)])
    wk = wk.reshape((1, d))
    t = torch.arange(n).reshape((n, 1))
    embedding[:, ::2] = torch.sin(t * wk[:, ::2])
    embedding[:, 1::2] = torch.cos(t * wk[:, ::2])

    return embedding


class MyUNet(nn.Module):
    def __init__(self, n_steps=1000, time_emb_dim=100):
        super(MyUNet, self).__init__()

        # Sinusoidal embedding
        self.time_embed = nn.Embedding(n_steps, time_emb_dim)
        self.time_embed.weight.data = sinusoidal_embedding(n_steps, time_emb_dim)
        self.time_embed.requires_grad_(False)

        # First half
        self.te1 = self._make_te(time_emb_dim, 1)
        self.b1 = nn.Sequential(
            MyBlock((1, 28, 28), 1, 10),
            MyBlock((10, 28, 28), 10, 10),
            MyBlock((10, 28, 28), 10, 10),
        )
        self.down1 = nn.Conv2d(10, 10, 4, 2, 1)

        self.te2 = self._make_te(time_emb_dim, 10)
        self.b2 = nn.Sequential(
            MyBlock((10, 14, 14), 10, 20),
            MyBlock((20, 14, 14), 20, 20),
            MyBlock((20, 14, 14), 20, 20),
        )
        self.down2 = nn.Conv2d(20, 20, 4, 2, 1)

        self.te3 = self._make_te(time_emb_dim, 20)
        self.b3 = nn.Sequential(
            MyBlock((20, 7, 7), 20, 40),
            MyBlock((40, 7, 7), 40, 40),
            MyBlock((40, 7, 7), 40, 40),
        )
        self.down3 = nn.Sequential(
            nn.Conv2d(40, 40, 2, 1), nn.SiLU(), nn.Conv2d(40, 40, 4, 2, 1)
        )

        # Bottleneck
        self.te_mid = self._make_te(time_emb_dim, 40)
        self.b_mid = nn.Sequential(
            MyBlock((40, 3, 3), 40, 20),
            MyBlock((20, 3, 3), 20, 20),
            MyBlock((20, 3, 3), 20, 40),
        )

        # Second half
        self.up1 = nn.Sequential(
            nn.ConvTranspose2d(40, 40, 4, 2, 1),
            nn.SiLU(),
            nn.ConvTranspose2d(40, 40, 2, 1),
        )

        self.te4 = self._make_te(time_emb_dim, 80)
        self.b4 = nn.Sequential(
            MyBlock((80, 7, 7), 80, 40),
            MyBlock((40, 7, 7), 40, 20),
            MyBlock((20, 7, 7), 20, 20),
        )

        self.up2 = nn.ConvTranspose2d(20, 20, 4, 2, 1)
        self.te5 = self._make_te(time_emb_dim, 40)
        self.b5 = nn.Sequential(
            MyBlock((40, 14, 14), 40, 20),
            MyBlock((20, 14, 14), 20, 10),
            MyBlock((10, 14, 14), 10, 10),
        )

        self.up3 = nn.ConvTranspose2d(10, 10, 4, 2, 1)
        self.te_out = self._make_te(time_emb_dim, 20)
        self.b_out = nn.Sequential(
            MyBlock((20, 28, 28), 20, 10),
            MyBlock((10, 28, 28), 10, 10),
            MyBlock((10, 28, 28), 10, 10, normalize=False),
        )

        self.conv_out = nn.Conv2d(10, 1, 3, 1, 1)

    def forward(self, x, t):
        # x is (N, 2, 28, 28) (image with positional embedding stacked on channel dimension)
        x = x.reshape([len(x), 1] + list(x.shape)[1:])
        t = self.time_embed((t * 999).int())  # TODO: should be 1000 not 999
        n = len(x)
        out1 = self.b1(x + self.te1(t).reshape(n, -1, 1, 1))  # (N, 10, 28, 28)
        out2 = self.b2(
            self.down1(out1) + self.te2(t).reshape(n, -1, 1, 1)
        )  # (N, 20, 14, 14)
        out3 = self.b3(
            self.down2(out2) + self.te3(t).reshape(n, -1, 1, 1)
        )  # (N, 40, 7, 7)

        out_mid = self.b_mid(
            self.down3(out3) + self.te_mid(t).reshape(n, -1, 1, 1)
        )  # (N, 40, 3, 3)

        out4 = torch.cat((out3, self.up1(out_mid)), dim=1)  # (N, 80, 7, 7)
        out4 = self.b4(out4 + self.te4(t).reshape(n, -1, 1, 1))  # (N, 20, 7, 7)

        out5 = torch.cat((out2, self.up2(out4)), dim=1)  # (N, 40, 14, 14)
        out5 = self.b5(out5 + self.te5(t).reshape(n, -1, 1, 1))  # (N, 10, 14, 14)

        out = torch.cat((out1, self.up3(out5)), dim=1)  # (N, 20, 28, 28)
        out = self.b_out(out + self.te_out(t).reshape(n, -1, 1, 1))  # (N, 1, 28, 28)

        out = self.conv_out(out)

        return out.squeeze()

    def _make_te(self, dim_in, dim_out):
        return nn.Sequential(
            nn.Linear(dim_in, dim_out), nn.SiLU(), nn.Linear(dim_out, dim_out)
        )

src/models/score_based/nn/unet.py
DavideScassola/score-based-constrained-generation
from pathlib import Path
from typing import Callable

import matplotlib.pyplot as plt
import numpy as np
import torch

from src.constants import ALLOW_LANGEVIN_CLEANING, LANGEVIN_CLEANING_PATIENCE
from src.models.model import Model
from src.models.score_based.nn.optimization import Optimization
from src.models.score_based.score_function import ScoreFunction
from src.models.score_based.score_matching import denoising_score_matching_loss
from src.report import score_plot
from src.util import load_json, store_json

from .score_matching import loss_plot, score_matching
from .sdes.sampler import SdeSolver, langevin_cleaning
from .sdes.sde import Sde


class ScoreBasedSde(Model):
    def __init__(
        self,
        *,
        sde: Sde,
        score_function_class,
        score_function_hyperparameters: dict,
        score_matching_loss: Callable = denoising_score_matching_loss,
        optimization: Optimization,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)

        self.sde = sde
        self.score_function_class = score_function_class
        self.score_function_hyperparameters = score_function_hyperparameters
        self.optimization = optimization
        self.score_loss_function = score_matching_loss
        self.train_losses = None

    def _train(self, X: torch.Tensor) -> None:
        self.shape = X.data[0].shape
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        self.score_model = self.score_function_class(
            shape=self.shape,
            sde=self.sde,
            device=self.device,
            **self.score_function_hyperparameters,
        )

        self.train_losses = score_matching(
            train_set=X,
            score_model=self.score_model,
            optimization=self.optimization,
            sm_loss=self.score_loss_function,
        )

    def _generate(
        self, n_samples: int, sde_solver: SdeSolver, **kwargs
    ) -> torch.Tensor:
        x = sde_solver.sample(
            sde=self.sde, score_function=self.score_model, n_samples=n_samples, **kwargs
        )

        constraint = getattr(self.score_model, "constraint", None)
        if constraint and ALLOW_LANGEVIN_CLEANING:
            langevin_cleaning(
                x=x,
                t=0.0,
                score_function=self.score_model,
                evaluator=self.score_model.constraint,
                patience=LANGEVIN_CLEANING_PATIENCE,
            )

        return x

    def _store(self, model_path: str) -> None:
        store_json(
            {"shape": self.shape, "device": str(self.device)},
            file=self.params_file(model_path),
        )
        self.score_model.store(model_path)

    def _load_(self, model_path: str) -> None:
        params = load_json(self.params_file(model_path))
        self.shape = params["shape"]
        self.device = torch.device(params["device"])
        self.score_model = self.score_function_class(
            shape=self.shape,
            sde=self.sde,
            device=self.device,
            **self.score_function_hyperparameters,
        )
        self.score_model.load_(model_path)

    def specific_report_plots(self, path: Path) -> None:
        if self.shape[0] == 1:
            score_plot(score_function=self.score_model, path=path)
        if self.train_losses:
            loss_plot(losses=self.train_losses, path=str(object=path))

src/models/score_based/score_based_sde.py
DavideScassola/score-based-constrained-generation
from abc import ABC, abstractmethod
from typing import List

import numpy as np
import torch
from torch import Tensor

from .nn.gine_mlp import MLP as GMLP
from .nn.time_residual_mlp import TimeResidualMLP
from .nn.unet import MyUNet
from .sdes.sde import VE, Sde

WEIGHTS_FILE = "model_weights"


class ScoreFunction(ABC):
    def __init__(
        self,
        *,
        sde: Sde,
        shape: tuple,
        device: str,
        rescale_factor_limit: float | None = 1000.0,
        add_prior_score_bias: bool = False,
        **hyperparameters,
    ) -> None:
        self.shape = shape
        self.sde = sde
        self.device = device
        self.rescale_factor_limit = rescale_factor_limit
        self.add_prior_score_bias = add_prior_score_bias
        self.model = self.build_model(**hyperparameters)
        # print(self.model)

    @abstractmethod
    def build_model(self, **hyperparameters) -> torch.nn.Module:
        pass

    @abstractmethod
    def forward(self, *, X: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        pass

    def __call__(self, *, X: torch.Tensor, t: torch.Tensor | float) -> torch.Tensor:
        with torch.no_grad():
            self.train(False)
            if isinstance(t, float):
                t = torch.full(
                    size=(len(X), 1), fill_value=t, dtype=torch.float32, device=X.device
                )
            return self.forward(X=X, t=t)

    def get_shape(self) -> tuple:
        return self.shape

    def parameters(self):
        return self.model.parameters()

    def train(self, train: bool = True):
        self.model.train(train)

    def input_scaler(self, *, X: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        if isinstance(self.sde, VE):
            _, std = self.sde.transition_kernel_mean_std(X=X, t=t)
            return X / torch.sqrt(1.0 + torch.square(std))
        return X

    def store(self, path: str):
        # TODO: using safe tensors would be better
        torch.save(self.model.state_dict(), f"{path}/{WEIGHTS_FILE}.pth")

    def load_(self, path: str):
        # TODO: using safe tensors would be better
        self.model.load_state_dict(torch.load(f"{path}/{WEIGHTS_FILE}.pth"))

    def weighted_score_coefficient(self, *, X, t: torch.Tensor) -> Tensor | float:
        # experimental
        """
        w(t) is defined such that:
        gaussian_score(t=0) * w(t) + (1-w(t)) * gaussian_score(t=1) = gaussian_score(t)

        so that at t = 1, w(t) = 0 and the score starts from the prior
        """
        if isinstance(self.sde, VE):
            mv = self.sde.sigma_max**2
            v = self.sde.transition_kernel_mean_std(X=X, t=t)[1] ** 2
            return (1 - v / mv) / (1 + v)
        else:
            return 1.0

    def prior_score_weighting(self, *, X, t: Tensor, score: Tensor):
        # experimental
        w = self.weighted_score_coefficient(X=X, t=t)
        return w * score + (1 - w) * self.prior_normal_score(X=X)

    def output_scaler(
        self, *, Y: torch.Tensor, X: torch.Tensor, t: torch.Tensor
    ) -> torch.Tensor:
        if not self.rescale_factor_limit:
            return Y
        _, std = self.sde.transition_kernel_mean_std(X=X, t=t)
        return Y * torch.min(
            1.0 / std,
            torch.full_like(input=std, fill_value=self.rescale_factor_limit),
        )

    def prior_normal_score(self, *, X: torch.Tensor) -> torch.Tensor:
        # TODO: one could generalize the formula for any sde
        return -X / (1 + self.sde.sigma_max**2) if isinstance(self.sde, VE) else -X


class MLP(ScoreFunction):
    def build_model(self, **hyperparameters) -> torch.nn.Module:
        numel = int(np.prod(np.array(self.get_shape())))
        return TimeResidualMLP(
            in_channels=numel, out_channels=numel, **hyperparameters
        ).to(self.device)

    def forward(self, *, X: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        scaled_X = self.input_scaler(X=X.flatten(1), t=t)
        output = (self.model.forward(X=scaled_X, t=t)).reshape([-1] + list(self.shape))
        nn_score = self.output_scaler(X=X, Y=output, t=t)
        if self.add_prior_score_bias:
            return nn_score + self.prior_normal_score(X=X)
        return nn_score


class GineMLP(ScoreFunction):
    def build_model(self, **hyperparameters) -> torch.nn.Module:
        numel = int(np.prod(np.array(self.get_shape())))
        return GMLP(n_features=numel, **hyperparameters).to(self.device)

    def forward(self, *, X: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        scaled_X = self.input_scaler(X=X.flatten(1), t=t)
        output = (self.model.forward(x=scaled_X, t=t)).reshape([-1] + list(self.shape))
        return output  # self.output_scaler(X=X, Y=output, t=t)


class Unet(ScoreFunction):
    def build_model(self, **hyperparameters) -> torch.nn.Module:
        return MyUNet(**hyperparameters).to(self.device)

    def forward(self, *, X: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        scaled_X = self.input_scaler(X=X, t=t)
        output = self.model.forward(scaled_X, t)  # - scaled_X
        return self.output_scaler(X=X, Y=output, t=t)

src/models/score_based/score_function.py
DavideScassola/score-based-constrained-generation
from functools import partial
from typing import Callable

import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import Tensor
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm

from src.constants import MINIMUM_SAMPLING_T
from src.models.score_based.nn.optimization import Optimization
from src.models.score_based.nn.parameters_ema import ParametersEMA
from src.util import reciprocal_distribution_sampler, uniform_sampler

from .score_function import ScoreFunction
from .sdes.sde import Sde

EPS = MINIMUM_SAMPLING_T * 0.1


def log_statistic(*, name: str, value) -> None:
    print(f"{name}: {value}")


def loss_plot(
    *, losses, epoch: int = 0, freq: int = 1, path: str | None = None
) -> None:
    loss_plot_name = "losses.png"
    file = f"{path}/{loss_plot_name}" if path else loss_plot_name
    if epoch % freq == 0:
        plt.plot(losses, alpha=0.8)
        plt.savefig(file)
        plt.close()


def denoising_score_matching_loss(
    *,
    X: torch.Tensor,
    score_model: ScoreFunction,
    train: bool = True,
    likelihood_weighting: bool = False,
    reciprocal_distribution_t: bool = False,
) -> torch.Tensor:
    # TODO: maybe it would be better to do this only one time, not every time the loss is computed
    score_model.train(train)
    t_size = (len(X), 1)

    t_sampler = (
        reciprocal_distribution_sampler
        if reciprocal_distribution_t
        else uniform_sampler
    )
    t = t_sampler(low=EPS, up=1.0, size=t_size, dtype=X.dtype, device=X.device)

    noise = torch.randn_like(X, device=X.device)
    mean, std = score_model.sde.transition_kernel_mean_std(X=X, t=t)
    X_t = mean + std * noise

    estimated_score = (
        score_model.forward(X=X_t, t=t) if train else score_model(X=X_t, t=t)
    )

    if not likelihood_weighting:
        losses = torch.square(estimated_score * std + noise)
    else:
        losses = torch.square(estimated_score + noise / std) * (
            score_model.sde.g(t=t) ** 2
        )
    return torch.mean(losses)


def accumulated_denoising_score_matching_loss(*, n: int = 20, **loss_args) -> Tensor:
    loss = denoising_score_matching_loss(**loss_args)
    for i in range(1, n):
        loss += denoising_score_matching_loss(**loss_args)
    return loss / n


def sliced_score_matching_loss(
    *,
    X: torch.Tensor,
    score_model: ScoreFunction,
    n_particles: int = 1,
    train: bool = True,
    noise=False,
) -> torch.Tensor:
    """
    From https://github.com/ermongroup/sliced_score_matching
    """
    # TODO: what's n_particles ?
    # TODO: maybe it would be better to do this only one time, not every time the loss is computed
    score_model.train(train)

    if noise:
        t = reciprocal_distribution_sampler(
            low=EPS, up=1.0, size=(len(X), 1), dtype=X.dtype, device=X.device
        )
        noise = torch.randn_like(X, device=X.device)
        mean, std = score_model.sde.transition_kernel_mean_std(X=X, t=t)
        X = mean + std * noise
    else:
        t = torch.zeros(size=(len(X), 1), device=X.device, dtype=X.dtype)

    dup_samples = (
        X.unsqueeze(0).expand(n_particles, *X.shape).contiguous().view(-1, *X.shape[1:])
    )
    dup_samples.requires_grad_(True)
    vectors = torch.randn_like(dup_samples)

    grad1 = score_model.forward(X=dup_samples, t=t)
    gradv = torch.sum(grad1 * vectors)
    loss1 = torch.sum(grad1 * grad1, dim=-1) / 2.0
    grad2 = torch.autograd.grad(gradv, dup_samples, create_graph=True)[0]
    loss2 = torch.sum(vectors * grad2, dim=-1)

    loss1 = loss1.view(n_particles, -1).mean(dim=0)
    loss2 = loss2.view(n_particles, -1).mean(dim=0)

    loss = loss1 + loss2
    return loss.mean()


def hybrid_ssm_dsm_loss(*, w: float, **kwargs):
    return w * sliced_score_matching_loss(noise=True, **kwargs) + (
        1 - w
    ) * denoising_score_matching_loss(reciprocal_distribution_t=False, **kwargs)


def score_matching(
    *,
    train_set: torch.Tensor,
    score_model: ScoreFunction,
    optimization: Optimization,
    sm_loss: Callable = denoising_score_matching_loss,
) -> list:
    tensor_train_set = TensorDataset(train_set.to(score_model.device))

    train_loader = DataLoader(
        dataset=tensor_train_set,
        batch_size=optimization.batch_size,
        shuffle=True,
        drop_last=True,
    )
    score_model.train(True)
    optimizer = optimization.build_optimizer(score_model.parameters())

    epochs_losses = []
    n_batches = len(train_loader)

    parameters_ema = (
        ParametersEMA(
            parameters=score_model.parameters(), decay=optimization.parameters_momentum
        )
        if optimization.parameters_momentum
        else None
    )

    for epoch in range(optimization.epochs):
        log_statistic(name="epoch", value=epoch + 1)
        losses = np.zeros(n_batches)

        pbar = tqdm(total=n_batches)
        for i, X_batch in enumerate(train_loader):
            optimizer.zero_grad()
            loss = sm_loss(X=X_batch[0], score_model=score_model, train=True)
            loss.backward()

            if optimization.gradient_clipping:
                torch.nn.utils.clip_grad_norm_(score_model.parameters(), max_norm=optimization.gradient_clipping)  # type: ignore

            optimizer.step()
            losses[i] = loss.detach().item()
            pbar.update(1)
        pbar.close()

        epoch_loss = np.mean(losses)
        epochs_losses.append(epoch_loss)
        log_statistic(name="train loss", value=epoch_loss)
        print()

        loss_plot(
            losses=epochs_losses,
            epoch=epoch,
            freq=1 + optimization.epochs // 5,
        )
        if parameters_ema:
            parameters_ema.update(score_model.parameters())

    if parameters_ema:
        parameters_ema.copy_to(score_model.parameters())
    loss_plot(losses=epochs_losses)

    return epochs_losses

src/models/score_based/score_matching.py
DavideScassola/score-based-constrained-generation
from typing import Callable

import numpy as np
import torch
from tqdm import tqdm

from src.constants import MINIMUM_SAMPLING_T

from ..score_function import ScoreFunction
from .sde import Sde

# TODO: try https://github.com/rtqichen/torchdiffeq or https://github.com/google-research/torchsde
# TODO: try scipy ode solvers
# TODO: try https://www.wikiwand.com/en/Stochastic_differential_equation


def langevin_update(*, x: torch.Tensor, score: torch.Tensor):
    # TODO: it doesn't work properly
    alpha = 1  # TODO: this works only for VE sde
    target_snr = 0.16  # TODO: another hyperparameter that should be settable
    noise = torch.randn_like(x)
    grad_norm = torch.norm(score.reshape(score.shape[0], -1), dim=-1).mean()
    noise_norm = torch.norm(noise.reshape(noise.shape[0], -1), dim=-1).mean()
    eps = (target_snr * noise_norm / grad_norm) ** 2 * 2 * alpha
    x += eps * score + noise * torch.sqrt(2.0 * eps)


def langevin_dynamics(
    *, x: torch.Tensor, t: float, score_function: ScoreFunction, steps: int
):
    it = range(steps) if steps < 1000 else tqdm(range(steps), desc="Langevyn cleaning")
    for _ in it:
        langevin_update(x=x, score=score_function(X=x, t=t))


def langevin_cleaning(
    *,
    x: torch.Tensor,
    t: float,
    score_function: ScoreFunction,
    patience: int,
    evaluator: Callable,
) -> None:
    waiting_time = 0
    v_max = evaluator(x=x).mean().item()

    pbar = tqdm(total=patience, desc="Langevyn cleaning...")
    while waiting_time < patience:
        langevin_update(x=x, score=score_function(X=x, t=t))
        v = evaluator(x=x).mean().item()
        if v > v_max:
            v_max = v
            waiting_time = 0
            pbar.reset()
        else:
            waiting_time = waiting_time + 1
            pbar.update()
    pbar.close()


class SdeSolver:
    def dx(
        self, *, x: torch.Tensor, sde: Sde, score: torch.Tensor, t: float, dt: float
    ) -> torch.Tensor | None:
        NotImplementedError("This solver doesn't explicitly define a formula for dx")

    def update(
        self, *, x: torch.Tensor, sde: Sde, score: torch.Tensor, t: float, dt: float
    ):
        x += self.dx(sde=sde, score=score, t=t, x=x, dt=dt)

    def solve(
        self,
        *,
        sde: Sde,
        score_function: ScoreFunction,
        X_0: torch.Tensor,
        steps: int = 1000,
        T: float = 1.0,
        corrector_steps: int = 0,
        final_corrector_steps: int = 0,
    ) -> torch.Tensor:
        x = X_0.clone()
        dt = T / steps
        for t in tqdm(
            np.linspace(T, MINIMUM_SAMPLING_T, steps), desc="Reversing the SDE"
        ):
            langevin_dynamics(
                x=x, t=t, score_function=score_function, steps=corrector_steps
            )
            self.update(x=x, sde=sde, score=score_function(X=x, t=t), t=t, dt=dt)
        if final_corrector_steps:
            langevin_dynamics(
                x=x,
                t=MINIMUM_SAMPLING_T,
                score_function=score_function,
                steps=final_corrector_steps,
            )
        return x

    def sample(
        self,
        *,
        sde: Sde,
        score_function: ScoreFunction,
        steps: int = 1000,
        T: float = 1.0,
        n_samples: int,
        corrector_steps: int = 0,
        final_corrector_steps: int = 0,
    ) -> torch.Tensor:
        X_0 = sde.prior_sampling(
            shape=score_function.get_shape(), n_samples=n_samples
        ).to(score_function.device)
        return self.solve(
            X_0=X_0,
            score_function=score_function,
            sde=sde,
            T=T,
            steps=steps,
            corrector_steps=corrector_steps,
            final_corrector_steps=final_corrector_steps,
        )


def dw(*, dt: float, shape: tuple, device):
    return torch.randn(size=shape, dtype=torch.float32, device=device) * (dt**0.5)


class EulerMethod(SdeSolver):
    def dx(
        self, *, x: torch.Tensor, sde: Sde, score: torch.Tensor, t: float, dt: float
    ) -> torch.Tensor:
        return -sde.reverse_f(X=x, t=t, score=score) * dt + sde.g(t=t) * dw(
            dt=dt, shape=x.shape, device=x.device
        )


class TimeAdaptiveEuler(EulerMethod):
    def __init__(self, *, dx_max: float = 0.01, dt_max: float = 0.01) -> None:
        super().__init__()
        self.dx_max = dx_max
        self.dt_max = dt_max

    def adaptive_dt(self, *, gradient: torch.Tensor) -> float:
        return self.dx_max / torch.max(gradient).item()

    def solve(
        self,
        *,
        sde: Sde,
        score_function: ScoreFunction,
        X_0: torch.Tensor,
        T: float = 1.0,
        corrector_steps: int = 0,
    ) -> torch.Tensor:
        x = X_0.clone()

        print("Reversing the SDE...")

        t = float(T)
        pbar = tqdm(total=T)
        while t > MINIMUM_SAMPLING_T:
            # print(f"{t/T:.1%}")
            score = score_function(X=x, t=t)
            dt = min(self.dt_max, self.adaptive_dt(gradient=score))

            if t - dt < MINIMUM_SAMPLING_T:
                dt = t - MINIMUM_SAMPLING_T

            for _ in range(corrector_steps):
                langevin_update(x=x, score=score)
            self.update(x=x, sde=sde, score=score, t=t, dt=dt)

            t -= dt
            pbar.update(dt)

        pbar.close()

        return x

    def sample(
        self,
        *,
        sde: Sde,
        score_function: ScoreFunction,
        dx_max: float = 0.01,
        dt_max: float = 0.01,
        T: float = 1.0,
        n_samples: int,
        corrector_steps: int = 0,
    ) -> torch.Tensor:
        X_0 = sde.prior_sampling(
            shape=score_function.get_shape(), n_samples=n_samples
        ).to(score_function.device)
        return self.solve(
            X_0=X_0,
            score_function=score_function,
            sde=sde,
            T=T,
            corrector_steps=corrector_steps,
        )

src/models/score_based/sdes/sampler.py
DavideScassola/score-based-constrained-generation
from abc import ABC, abstractmethod
from typing import Tuple

import numpy as np
import torch
from torch import Tensor


def expand_like(input: Tensor, *, like: Tensor):
    return (
        input.reshape(tuple(input.shape) + (1,) * (len(like.shape) - len(input.shape)))
        .expand_as(like)
        .to(like.device)
    )


class Sde(ABC):
    # TODO: in the paper T is a property of the SDE. T is probably such that
    # X_T ~ prior, but why not fixing T=1 and defining the sde such that x_t=1
    # ~ pior ?

    @abstractmethod
    def f(self, *, X: Tensor, t: float | Tensor) -> Tensor:
        """
        drift coefficient f(x,t)
        """
        pass

    @abstractmethod
    def g(self, *, t: float | Tensor) -> float | Tensor:
        """
        diffusion coefficient g(t)
        """
        pass

    @abstractmethod
    def prior_sampling(self, *, shape: tuple, n_samples: int) -> Tensor:
        pass

    @abstractmethod
    def snr(self, t: float) -> Tensor:
        pass

    @abstractmethod
    def denoising(self, *, t: Tensor, score: Tensor, X: Tensor) -> Tensor:
        pass

    def transition_kernel_mean_std(
        self, *, X: Tensor, t: Tensor
    ) -> Tuple[Tensor, Tensor]:
        raise NotImplementedError(
            "This sde has not an explicit gaussian transition kernel"
        )

    def reverse_f(self, *, X: Tensor, t: float, score: Tensor) -> Tensor:
        return self.f(X=X, t=t) - ((self.g(t=t)) ** 2) * score


class VE(Sde):
    def __init__(self, *, sigma_min=0.01, sigma_max=50.0) -> None:
        self.sigma_min = sigma_min
        self.sigma_max = sigma_max

    def f(self, *, X: Tensor, t: float | Tensor) -> Tensor:
        return torch.zeros_like(X)

    def g(self, *, t: float | Tensor) -> float | Tensor:
        sigma = self.sigma_min * (self.sigma_max / self.sigma_min) ** t
        return sigma * np.sqrt(2 * (np.log(self.sigma_max) - np.log(self.sigma_min)))

    def transition_kernel_mean_std(
        self, *, X: Tensor, t: Tensor
    ) -> Tuple[Tensor, Tensor]:
        return X, expand_like(
            self.sigma_min * (self.sigma_max / self.sigma_min) ** t, like=X
        )

    def prior_sampling(self, shape: tuple, n_samples: int) -> Tensor:
        return torch.randn(size=[n_samples] + list(shape)) * self.sigma_max

    def snr(self, t: float) -> Tensor:
        std = self.transition_kernel_mean_std(X=torch.tensor(0.0), t=torch.tensor(t))[1]
        return (1 + std**2) ** (-0.5)

    def denoising(self, *, t: Tensor, score: Tensor, X: Tensor) -> Tensor:
        std = self.transition_kernel_mean_std(X=X, t=t)[1]
        return X + score * (std**2)


class subVP(Sde):
    def __init__(self, *, beta_min: float = 0.1, beta_max: float = 20.0) -> None:
        self.beta_0 = beta_min
        self.beta_1 = beta_max

    def f(self, *, X: Tensor, t: float | Tensor) -> Tensor:
        beta_t = self.beta_0 + t * (self.beta_1 - self.beta_0)
        return -0.5 * beta_t * X

    def g(self, *, t: float | Tensor) -> float | Tensor:
        beta_t = self.beta_0 + t * (self.beta_1 - self.beta_0)

        if isinstance(t, Tensor):
            discount = 1.0 - torch.exp(
                -2.0 * self.beta_0 * t - (self.beta_1 - self.beta_0) * t**2
            )
            return torch.sqrt(beta_t * discount)
        else:
            discount = 1.0 - np.exp(
                -2.0 * self.beta_0 * t - (self.beta_1 - self.beta_0) * t**2
            )
            return np.sqrt(beta_t * discount)

    def transition_kernel_mean_std(
        self, *, X: Tensor, t: Tensor
    ) -> Tuple[Tensor, Tensor]:
        log_mean_coeff = expand_like(
            -0.25 * t**2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0, like=X
        )
        mean = torch.exp(log_mean_coeff) * X
        std = 1 - torch.exp(2.0 * log_mean_coeff)
        return mean, std

    def prior_sampling(self, shape: tuple, n_samples: int) -> Tensor:
        return torch.randn(size=[n_samples] + list(shape))

    def snr(self, t: float) -> Tensor:
        log_mean_coeff = torch.tensor(
            -0.25 * t**2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
        )
        k = torch.exp(2 * log_mean_coeff)
        std = 1 - k
        return k / (k**2 + std**2) ** 0.5

    def denoising(self, *, t: Tensor, score: Tensor, X: Tensor) -> Tensor:
        log_mean_coeff = expand_like(
            -0.25 * t**2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0, like=X
        )
        std = 1 - torch.exp(2.0 * log_mean_coeff)
        return (X + score * (std**2)) / torch.exp(log_mean_coeff)

src/models/score_based/sdes/sde.py
DavideScassola/score-based-constrained-generation
import os

from src.data import Dataset
from src.models.tabular_model import TabularModel
from src.preprocessors.to_tensor import ToTensor
from src.report import *
from src.util import rejection_sampling, store_json

from .score_based_sde import ScoreBasedSde


def constraint_rejection_sampling(df: pd.DataFrame, constraint: Constraint):
    probability_distribution = lambda x: np.exp(
        constraint.strength * constraint.f(torch.from_numpy(x)).numpy()
    )
    x = rejection_sampling(df.to_numpy(), probability_distribution)
    return pd.DataFrame(x, columns=df.columns)


def refill_rejection_sampling(
    df: pd.DataFrame, constraint: Constraint, sampler, generation_options
):
    df_rejection_sampling = constraint_rejection_sampling(df, constraint)

    discarded_samples = len(df) - len(df_rejection_sampling)
    # TODO: make this batched
    if discarded_samples > 0:
        acceptance_rate = len(df_rejection_sampling) / len(df)

        generation_options["n_samples"] = int(
            1.2 * discarded_samples // acceptance_rate
        )

        print("not enough samples after rejection, generating again ...")
        df_new = sampler(**generation_options)
        df_rejection_sampling_new = constraint_rejection_sampling(df_new, constraint)
        df_rejection_sampling = pd.concat(
            [df_rejection_sampling, df_rejection_sampling_new]
        )

    return df_rejection_sampling.iloc[: len(df)], len(df_rejection_sampling) / (
        len(df_new) + len(df)
    )


SHOW_REJECTION_SAMPLING = True


class TabularScoreBasedSde(ScoreBasedSde, TabularModel):
    def generate_report(
        self,
        *,
        path: str | Path,
        dataset: Dataset,
        generation_options: dict,
        constraint: Constraint | None = None
    ):
        report_folder = path / Path(REPORT_FOLDER_NAME)
        os.makedirs(report_folder, exist_ok=False)

        samples = self.generate(**generation_options)
        store_samples(df_generated=samples, path=report_folder, name="guidance.csv")

        if constraint:
            # TODO: setting the score to None temporarily is not so elegant
            original_score = self.score_model.constraint
            self.score_model.constraint = None
            df_train = self.generate(**generation_options)
            store_samples(
                df_generated=df_train, path=report_folder, name="no_guidance.csv"
            )

            if SHOW_REJECTION_SAMPLING:
                df_train, acceptance_rate = refill_rejection_sampling(
                    df_train, constraint, self.generate, generation_options
                )
                store_json(
                    {"acceptance_rate": acceptance_rate},
                    file=report_folder / Path("acceptance_rate.json"),
                )
                self.score_model.constraint = original_score
                store_samples(
                    df_generated=df_train,
                    path=report_folder,
                    name="rejection_sampling.csv",
                )

            if len(samples.columns) > 1:
                # TODO: exp?
                constraint_exp = lambda x: np.exp(
                    constraint.strength
                    * constraint.f(ToTensor().fit_transform(x)).numpy()
                )
                df_train["constraint_value"] = constraint_exp(df_train)
                samples["constraint_value"] = constraint_exp(samples)
                rcParams.update({"figure.autolayout": True})
        else:
            df_train = dataset.get(train=True)

        print("plotting...")
        for comparison in (
            histograms_comparison,
            statistics_comparison,
            correlations_comparison,
        ):
            comparison(
                df_generated=samples,
                name_generated="guidance" if constraint else "generated",
                df_train=df_train,
                name_original=(
                    "rejection sampling" if SHOW_REJECTION_SAMPLING else "no guidance"
                )
                if constraint
                else "true",
                path=report_folder,
                constraint=None if SHOW_REJECTION_SAMPLING else constraint,
                model=self,
            )

        self.specific_report_plots(report_folder)

src/models/score_based/tabular_score_based_sde.py
DavideScassola/score-based-constrained-generation
import os
from typing import Callable, Tuple

from torch import Tensor

from src.data import Dataset
from src.models.score_based.nn.optimization import Optimization
from src.models.score_based.score_matching import denoising_score_matching_loss
from src.models.score_based.sdes.sde import Sde
from src.models.time_series_model import TimeSeriesModel
from src.report import *
from src.util import store_json

from .score_based_sde import ScoreBasedSde


def store_array(x: np.ndarray | Tensor, path: Path, name: str) -> None:
    if isinstance(x, Tensor):
        x = x.cpu().numpy()
    np.save(f"{path}/{name}", x)


def torch_rejection_sampling(x: Tensor, acceptance_function) -> Tensor:
    """acceptance_function is a probability density function"""
    acceptance = acceptance_function(x)
    assert tuple(acceptance.shape) == (len(x),)

    max_acceptance = torch.max(acceptance)
    if max_acceptance.item() <= 0:
        print(f"Warning: max_acceptance is {max_acceptance}, no sample can be accepted")
        return Tensor()
    acceptance /= max_acceptance
    selection = torch.bernoulli(acceptance) == 1
    return x[selection]


def constraint_rejection_sampling(X: Tensor, constraint: Constraint) -> Tensor:
    probability_distribution = lambda x: torch.exp(
        constraint.strength * constraint.f(x)
    )
    return torch_rejection_sampling(X, probability_distribution)


def refill_rejection_sampling(
    X: torch.Tensor, constraint: Constraint, sampler, generation_options
) -> Tuple[Tensor, float]:
    X_rejection_sampling = constraint_rejection_sampling(X, constraint)

    discarded_samples = len(X) - len(X_rejection_sampling)
    if discarded_samples > 0:
        acceptance_rate = len(X_rejection_sampling) / len(X)

        if acceptance_rate == 0:
            print(f"Warning: all samples were rejected")
            return Tensor(), 0

        if (discarded_samples // acceptance_rate) > 1e4:
            print(
                f"Warning: cannot generate {discarded_samples // acceptance_rate} samples"
            )
            return Tensor(), 0

        generation_options["n_samples"] = int(discarded_samples // acceptance_rate)

        print("not enough samples after rejection, generating again ...")
        X_new = sampler(**generation_options)
        X_rejection_sampling_new = constraint_rejection_sampling(X_new, constraint)
        X_rejection_sampling = torch.concat(
            [X_rejection_sampling, X_rejection_sampling_new]
        )

    return X_rejection_sampling[: len(X)], len(X_rejection_sampling) / (
        len(X) + len(X_new)
    )


class TimeSeriesScoreBasedSde(ScoreBasedSde, TimeSeriesModel):
    def __init__(
        self,
        *,
        sde: Sde,
        score_function_class,
        score_function_hyperparameters: dict,
        optimization: Optimization,
        names: list[str],
        score_matching_loss: Callable = denoising_score_matching_loss,
        **kwargs,
    ) -> None:
        super().__init__(
            sde=sde,
            score_function_class=score_function_class,
            score_function_hyperparameters=score_function_hyperparameters,
            score_matching_loss=score_matching_loss,
            optimization=optimization,
            **kwargs,
        )
        self.names = names

    def generate_report(
        self,
        *,
        path: str | Path,
        dataset: Dataset,
        generation_options: dict,
        constraint: Constraint | None = None,
    ):
        report_folder = path / Path(REPORT_FOLDER_NAME)
        os.makedirs(report_folder, exist_ok=False)

        print("Sampling from constrained model...")
        samples = self.generate(**generation_options)
        plots_folder = str(path) + "/samples"
        os.makedirs(plots_folder)

        store_array(samples, path=report_folder, name="guidance.npy")

        print("Plotting samples from constrained model...")
        time_series_plots(
            samples=samples.cpu().numpy(),
            folder=plots_folder,
            n_plots=10,
            series_per_plot=10,
            features_names=self.names,
        )
        if constraint:
            satisfaction_plot(
                X=samples,
                constraint=constraint,
                path=report_folder,
                label="guidance",
                color="orange",
            )

        if not constraint:
            print("Plotting data from train set...")
            plots_folder = str(path) + "/samples_from_train"
            os.makedirs(plots_folder)

            time_series_plots(
                samples=np.random.permutation(dataset.get(train=True).cpu().numpy()),
                folder=plots_folder,
                n_plots=5,
                series_per_plot=15,
                features_names=self.names,
                color="blue",
            )
        else:
            print("Plotting samples from rejection sampling...")
            plots_folder = str(path) + "/rejection_sampling"
            os.makedirs(plots_folder)

            # TODO: setting the score to None temporarily is not so elegant
            original_score = self.score_model.constraint
            self.score_model.constraint = None

            print("Sampling from unconstrained model...")
            samples = self.generate(**generation_options)
            store_array(samples, path=report_folder, name="no_guidance.npy")

            samples, acceptance_rate = refill_rejection_sampling(
                samples, constraint, self.generate, generation_options
            )
            store_array(samples, path=report_folder, name="rejection_sampling.npy")

            store_json(
                {"acceptance_rate": acceptance_rate},
                file=report_folder / Path("acceptance_rate.json"),
            )
            self.score_model.constraint = original_score

            if len(samples) > 0:
                time_series_plots(
                    samples=samples.cpu().numpy(),
                    folder=plots_folder,
                    n_plots=10,
                    series_per_plot=10,
                    features_names=self.names,
                    color="green",
                )

                satisfaction_plot(
                    X=samples,
                    constraint=constraint,
                    path=report_folder,
                    label="rejection sampling",
                    color="green",
                )

src/models/score_based/time_series_score_based_sde.py
DavideScassola/score-based-constrained-generation
import os
from pathlib import Path
from typing import List

import pandas as pd
import torch

from src.data import Dataset
from src.preprocessors.preprocessor import (TensorPreprocessor,
                                            composed_fit_transform,
                                            composed_inverse_transform)
from src.preprocessors.to_tensor import ToTensor
from src.report import *

from .model import Model


class TabularModel(Model):
    def __init__(
        self, *, tensor_preprocessors: List[TensorPreprocessor] = [], numpy=False
    ) -> None:
        super().__init__()
        self.tensor_preprocessors = tensor_preprocessors
        self.df2tensor = ToTensor()
        self.preprocessors = [self.df2tensor] + self.tensor_preprocessors

    def train(self, dataset: Dataset):
        df: pd.DataFrame = dataset.get(train=True)
        X: torch.Tensor = composed_fit_transform(
            self.df2tensor.fit_transform(df), preprocessors=self.tensor_preprocessors
        )
        self._train(X)

    def generate(self, n_samples: int, **kwargs) -> pd.DataFrame:
        samples = self._generate(n_samples, **kwargs)
        return self.df2tensor.reverse_transform(
            composed_inverse_transform(samples, preprocessors=self.tensor_preprocessors)
        )

    def generate_report(
        self, *, path: str | Path, dataset: Dataset, generation_options: dict
    ):
        report_folder = path / Path(REPORT_FOLDER_NAME)
        os.makedirs(report_folder, exist_ok=False)

        samples = self.generate(**generation_options)

        df_train = dataset.get(train=True)

        store_samples(df_generated=samples, path=report_folder)

        print("plotting...")
        for comparison in (
            histograms_comparison,
            statistics_comparison,
            correlations_comparison,
        ):
            comparison(
                df_generated=samples,
                df_train=df_train,
                path=report_folder,
                model=self,
            )

        self.specific_report_plots(report_folder)

src/models/tabular_model.py
DavideScassola/score-based-constrained-generation
import os
from pathlib import Path
from typing import List

import torch

from src.data import Dataset
from src.preprocessors.preprocessor import (TensorPreprocessor,
                                            composed_fit_transform,
                                            composed_inverse_transform)
from src.report import *

from .model import Model


class TimeSeriesModel(Model):
    def __init__(self, *, tensor_preprocessors: List[TensorPreprocessor] = []) -> None:
        self.tensor_preprocessors = tensor_preprocessors
        self.preprocessors = self.tensor_preprocessors

    def train(self, dataset: Dataset):
        X = dataset.get(train=True)
        if not isinstance(X, torch.Tensor):
            raise ValueError(f"dataset {dataset} should be a tensor")
        X = composed_fit_transform(X, preprocessors=self.tensor_preprocessors)
        self._train(X)

    def generate(self, n_samples: int, **kwargs) -> torch.Tensor:
        samples = self._generate(n_samples, **kwargs)
        return composed_inverse_transform(
            samples, preprocessors=self.tensor_preprocessors
        )

    def generate_report(
        self, *, path: str | Path, dataset: Dataset, generation_options: dict
    ) -> None:
        raise NotImplementedError()

src/models/time_series_model.py
DavideScassola/score-based-constrained-generation
import torch

from .preprocessor import TensorPreprocessor


class Discretizer(TensorPreprocessor):
    # TODO: check gradient behavior

    def transform(self, x: torch.Tensor) -> torch.Tensor:
        return x

    def reverse_transform(self, x: torch.Tensor) -> torch.Tensor:
        return torch.round(x)

src/preprocessors/discretizer.py
DavideScassola/score-based-constrained-generation
import torch

from .preprocessor import TensorPreprocessor

EPS = 1e-6


class Log(TensorPreprocessor):
    def transform(self, x: torch.Tensor) -> torch.Tensor:
        return torch.log(x + EPS)

    def reverse_transform(self, x: torch.Tensor) -> torch.Tensor:
        return torch.exp(x) - EPS

src/preprocessors/log.py
DavideScassola/score-based-constrained-generation
import numpy as np
import torch

from .preprocessor import TensorPreprocessor


class MeanCovNormalizer(TensorPreprocessor):
    def fit(self, x: torch.Tensor):
        self.parameters["mean"] = np.mean(x, axis=0)
        self.parameters["costd"] = np.linalg.cholesky(np.cov(x.T)).T

    def transform(self, x: torch.Tensor) -> torch.Tensor:
        return (x - self.parameters["mean"]).dot(
            np.linalg.inv(self.parameters["costd"])
        )

    def reverse_transform(self, x: torch.Tensor) -> torch.Tensor:
        return x.dot(self.parameters["costd"]) + self.parameters["mean"]

    def serialize(self, p: dict) -> dict:
        # TODO
        return {k: v.tolist() for k, v in p.items()}

    def deserialize(self, p: dict) -> dict:
        # TODO
        return {k: torch.tensor(v) for k, v in p.items()}

src/preprocessors/mean_cov_normalizer.py
DavideScassola/score-based-constrained-generation
import torch

from .preprocessor import TensorPreprocessor


class MeanStdNormalizer(TensorPreprocessor):
    def fit(self, x: torch.Tensor):
        self.parameters["mean"] = x.mean(axis=0)  # type: ignore
        self.parameters["std"] = x.std(axis=0)  # type: ignore
        self.parameters["std"][
            self.parameters["std"] == 0.0
        ] = 1.0  # Avoiding 0. division

    def transform(self, x: torch.Tensor) -> torch.Tensor:
        self.parameters_to_device(x)
        return (x - self.parameters["mean"]) / self.parameters["std"]

    def parameters_to_device(self, x: torch.Tensor):
        if x.device != self.parameters["mean"].device:
            self.parameters["mean"] = self.parameters["mean"].to(x.device)
            self.parameters["std"] = self.parameters["std"].to(x.device)

    def reverse_transform(self, x: torch.Tensor) -> torch.Tensor:
        self.parameters_to_device(x)
        return x * self.parameters["std"] + self.parameters["mean"]

    def serialize(self, p: dict) -> dict:
        return {k: v.tolist() for k, v in p.items()}

    def deserialize(self, p: dict) -> dict:
        return {k: torch.tensor(v) for k, v in p.items()}

src/preprocessors/mean_std_normalizer.py
DavideScassola/score-based-constrained-generation
from abc import ABC, abstractmethod
from pathlib import Path
from typing import List

import torch

from src.util import load_json, store_json


class Preprocessor(ABC):
    def __init__(self):
        self.parameters = {}

    def fit(self, x):
        pass

    @abstractmethod
    def transform(self, x):
        pass

    @abstractmethod
    def reverse_transform(self, x):
        pass

    def fit_transform(self, x):
        self.fit(x)
        return self.transform(x)

    def parameters_file(self, model_path: str | Path, *, extension: str = ".json"):
        # TODO: this works only if every kind of preprocessor is used only once
        return Path(model_path) / Path(str(self.__class__.__name__) + extension)

    def load_(self, model_path: str | Path):
        self.parameters = self.deserialize(load_json(self.parameters_file(model_path)))

    def store(self, model_path: str | Path):
        store_json(
            self.serialize(self.parameters), file=self.parameters_file(model_path)
        )

    def serialize(self, p: dict):
        return p

    def deserialize(self, p: dict):
        return p


class TensorPreprocessor(Preprocessor, ABC):
    def fit(self, x: torch.Tensor) -> None:
        pass

    @abstractmethod
    def transform(self, x: torch.Tensor) -> torch.Tensor:
        pass

    @abstractmethod
    def reverse_transform(self, x: torch.Tensor) -> torch.Tensor:
        pass

    def fit_transform(self, x: torch.Tensor) -> torch.Tensor:
        self.fit(x)
        return self.transform(x)


def composed_fit_transform(
    input: torch.Tensor, *, preprocessors: List[TensorPreprocessor]
) -> torch.Tensor:
    if len(preprocessors) == 0:
        return input
    out = preprocessors[0].fit_transform(input)
    for p in preprocessors[1:]:
        out = p.fit_transform(out)
    return out


def composed_inverse_transform(
    input: torch.Tensor, *, preprocessors: List[TensorPreprocessor]
) -> torch.Tensor:
    if len(preprocessors) == 0:
        return input
    out = preprocessors[-1].reverse_transform(input)
    for p in preprocessors[-2::-1]:
        out = p.reverse_transform(out)
    return out

src/preprocessors/preprocessor.py
DavideScassola/score-based-constrained-generation
import numpy as np
from sklearn.preprocessing import QuantileTransformer

from .preprocessor import Preprocessor


class QuantileNormalizer(Preprocessor):
    # TODO: update this, and add non-differentiable transformers like this to model class
    def __init__(self, output_distribution: str = "normal"):
        self.transformer = QuantileTransformer(output_distribution=output_distribution)
        self.fit = self.transformer.fit
        self.transform = self.transformer.transform
        self.reverse_transform = self.transformer.inverse_transform

src/preprocessors/quantile_normalizer.py
DavideScassola/score-based-constrained-generation
import torch

from .preprocessor import TensorPreprocessor


class Rescaler(TensorPreprocessor):
    def __init__(
        self,
        *,
        subtract: float = 0.0,
        add: float = 0.0,
        divide_by: float = 1.0,
        multiply_by: float = 1.0
    ):
        super().__init__()
        self.parameters["subtract"] = subtract - add
        self.parameters["divide_by"] = divide_by
        self.parameters["multiply_by"] = multiply_by

    def fit(self, x: torch.Tensor):
        pass

    def transform(self, x: torch.Tensor) -> torch.Tensor:
        return (
            self.parameters["multiply_by"]
            * (x - self.parameters["subtract"])
            / self.parameters["divide_by"]
        )

    def reverse_transform(self, x: torch.Tensor) -> torch.Tensor:
        return (
            x * self.parameters["divide_by"] / self.parameters["multiply_by"]
            + self.parameters["subtract"]
        )

    def serialize(self, p: dict) -> dict:
        return {k: v for k, v in p.items()}

    def deserialize(self, p: dict) -> dict:
        return {k: v for k, v in p.items()}

src/preprocessors/rescaler.py
DavideScassola/score-based-constrained-generation
import numpy as np
import pandas as pd
from sklearn.preprocessing import LabelEncoder

from .preprocessor import Preprocessor


def is_categorical(s: pd.Series):
    # TODO: by default any 'object' is a category
    return s.dtype == "O"


class StringToInt(Preprocessor):
    def fit(self, df: pd.DataFrame):
        pass

    def transform(self, df: pd.DataFrame):
        pass

    def fit_transform(self, df: pd.DataFrame) -> pd.DataFrame:
        self.label_encoder = {
            c: LabelEncoder() for c in df.columns if is_categorical(df[c])
        }
        for v in self.label_encoder:
            df[v] = self.label_encoder[v].fit_transform(df[v])

        self.parameters["categories"] = {column_name: list(enc.classes_) for column_name, enc in self.label_encoder.items()}  # type: ignore

        # putting categories at the end
        self.parameters["names_order"] = list(df.columns)
        self.parameters["continuous_columns"] = [
            c
            for c in self.parameters["names_order"]
            if c not in self.parameters["categories"]
        ]
        self.parameters["categorical_columns"] = list(
            self.parameters["categories"].keys()
        )
        return df[
            self.parameters["continuous_columns"]
            + self.parameters["categorical_columns"]
        ]

    def reverse_transform(self, df_transformed: pd.DataFrame) -> pd.DataFrame:
        # restoring original columns order
        df = df_transformed[self.parameters["names_order"]]

        for categorical_column, classes in self.parameters["categories"].items():
            le = LabelEncoder()
            le.classes_ = np.array(classes)
            df[categorical_column] = le.inverse_transform(df[categorical_column])
        return df

src/preprocessors/string_to_int.py
DavideScassola/score-based-constrained-generation
import numpy as np
import pandas as pd

from .preprocessor import Preprocessor


class ToNumpy(Preprocessor):
    def __init__(self, float_type: type = np.float32):
        super().__init__()
        self.float_type = float_type

    def fit(self, x: pd.DataFrame):
        self.parameters["names"] = list(x.columns)

    def transform(self, x: pd.DataFrame):
        return x.to_numpy().astype(self.float_type) if self.float_type else x.to_numpy()

    def reverse_transform(self, x: np.ndarray):
        return pd.DataFrame(x, columns=self.parameters["names"])

src/preprocessors/to_numpy.py
DavideScassola/score-based-constrained-generation
import pandas as pd
import torch

from .preprocessor import Preprocessor


class ToTensor(Preprocessor):
    def fit(self, x: pd.DataFrame):
        self.parameters["names"] = list(x.columns)

    def transform(self, df: pd.DataFrame) -> torch.Tensor:
        return torch.from_numpy(df.values).float()

    def fit_transform(self, x: pd.DataFrame) -> torch.Tensor:
        self.fit(x)
        return self.transform(x)

    def reverse_transform(self, x: torch.Tensor):
        return pd.DataFrame(x.cpu().numpy(), columns=self.parameters["names"])

src/preprocessors/to_tensor.py
DavideScassola/score-based-constrained-generation
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import torch
import tqdm
from matplotlib import image, rcParams
from scipy import stats
from scipy.stats import kstest
from sklearn.neighbors import KernelDensity
from torch import Tensor

from src.constants import MINIMUM_SAMPLING_T
from src.constraints.constraint import Constraint
from src.util import l1_divergence, max_symmetric_D_kl, store_json, upscale

SAMPLES_NAME = "samples.csv"
CORRELATIONS_NAME = "correlations"
HISTOGRAMS_NAME = "histograms"
STATS_NAME = "stats.json"
REPORT_FOLDER_NAME = "report"
IMAGE_FORMAT = "png"


def kstest_pvalue_and_bool(x, y):
    result = kstest(x, y)
    return (1 if result.pvalue > 0.1 else 0, result.pvalue)


def list_corr(df):
    m = df.corr().to_numpy()
    return m[np.triu_indices_from(m, 1)]


def correlation_matrix_similarities(df1: pd.DataFrame, df2: pd.DataFrame):
    if "constraint_value" in df1.columns:
        df1 = df1.drop("constraint_value", axis=1)
        df2 = df2.drop("constraint_value", axis=1)
    c1 = list_corr(df1)
    c2 = list_corr(df2)

    return {
        # "pearson corr": stats.pearsonr(c1, c2),
        # "spearman corr": stats.spearmanr(c1, c2),
        "average l1": np.mean(abs(c1 - c2)),
    }


NUMERICAL_STATS = {
    "mean": lambda x: np.mean(x, axis=0),
    "std": lambda x: np.std(x, axis=0, ddof=1),
    "correlations": lambda df: df.corr(),
}

COLUMNWISE_NUMERICAL_COMPARISON_STATS = {
    "l1_divergence": l1_divergence,
    "max_symmetric_D_kl": max_symmetric_D_kl,
    "two-sample Kolmogorov-Smirnov test": kstest_pvalue_and_bool,
}

GLOBAL_NUMERICAL_COMPARISON_STATS = {
    "correlation matrix similarity": correlation_matrix_similarities
}

CATEGORICAL_STATS = {}

DEFAULT_KDE_ARGS = {}


def kde_plot(
    samples: pd.Series,
    *,
    ax=None,
    kde_args: dict = DEFAULT_KDE_ARGS,
    plt_args: dict = {},
    start=None,
    end=None,
    constraint: None | Constraint = None,
):
    ax = ax if ax else plt
    x = np.linspace(
        np.min(samples) if start is None else start,
        np.max(samples) if end is None else end,
        1000,
    )
    c = (
        constraint.f(torch.from_numpy(x)).numpy() * constraint.strength
        if constraint
        else 0.0
    )
    kde = KernelDensity(**kde_args).fit(samples.to_numpy().reshape(-1, 1))
    density = np.exp(kde.score_samples(x.reshape(-1, 1)) + c)
    density /= np.trapz(y=density, x=x)
    ax.plot(x, density, **plt_args)


def bins(s: pd.Series):
    return np.linspace(min(s), max(s), int(len(s) ** 0.5))


def constraint_plot(
    *, start: float, end: float, constraint: Constraint, max_y: float | None = None
):
    x = np.linspace(start, end, 1000)
    c = np.exp(constraint.f(torch.from_numpy(x)).numpy() * constraint.strength)
    if max_y:
        c -= np.max(c) * 1.05
        c *= 0.25 * max_y / abs(np.min(c))
    plt.plot(x, c, label="constraint", alpha=0.8, color="black", linestyle="--")


def histograms_comparison(
    *,
    df_generated: pd.DataFrame,
    df_train: pd.DataFrame,
    path: Path,
    name_generated="generated",
    name_original="original",
    **kwargs,
):
    constraint = kwargs["constraint"]

    n = len(df_generated.columns)
    fig, axes = plt.subplots(n, 1, sharex=False)
    if n == 1:
        axes = [axes]
    fig.set_size_inches(6, n + 4)
    fig.suptitle("Histogram Comparison")

    plt.tight_layout()

    for i, c in enumerate(df_generated.columns):
        is_categorical = df_generated[c].dtype in (str, "O")

        if is_categorical:
            common_args = dict(
                x=c,
                stat="proportion",
                alpha=0.5,
            )
        else:
            data = pd.Series(pd.concat((df_generated[c], df_train[c])))
            std = data.std()
            start = np.min(data.to_numpy()) - 0.5 * std
            end = np.max(data.to_numpy()) + 0.5 * std
            b = bins(pd.Series(pd.concat((df_generated[c], df_train[c]))))
            start = start
            end = end
            common_args = dict(x=c, stat="density", alpha=0.5, bins=b)

        ax_true = sns.histplot(
            ax=axes[i], data=df_train, label=name_original, **common_args
        )

        bandwidth = min(df_train[c].std(), df_generated[c].std()) / 10

        kde_plot(
            df_train[c],
            ax=axes[i],
            start=start,  # type: ignore
            end=end,  # type: ignore
            kde_args={"bandwidth": bandwidth},
            plt_args={"color": "blue", "alpha": 0.5, "linestyle": "-"},
        )

        ax_generated = sns.histplot(
            ax=axes[i], data=df_generated, label=name_generated, **common_args
        )
        kde_plot(
            df_generated[c],
            ax=axes[i],
            start=start,  # type: ignore
            end=end,  # type: ignore
            kde_args={"bandwidth": bandwidth},
            plt_args={"color": "orange", "alpha": 0.8, "linestyle": "-"},
        )

        if constraint and n == 1 and not is_categorical:
            constraint_plot(
                start=start,  # type: ignore
                end=end,  # type: ignore
                max_y=max(ax_true.get_ylim()[1], ax_generated.get_ylim()[1]),
                constraint=constraint,
            )

            kde_plot(
                df_train[c],
                start=start,  # type: ignore
                end=end,  # type: ignore
                kde_args={"bandwidth": bandwidth},
                plt_args={
                    "color": "blue",
                    "alpha": 0.8,
                    "linestyle": "--",
                    "label": "original constrained",
                },
                constraint=constraint,
            )

    plt.legend()
    plt.savefig(path / Path(f"{HISTOGRAMS_NAME}.{IMAGE_FORMAT}"))
    plt.close()


def store_samples(*, df_generated: pd.DataFrame, path: Path, name: str = SAMPLES_NAME):
    df_generated.to_csv(path / Path(name), index=False)


def get_stats(df: pd.DataFrame):
    df_numerical = df.select_dtypes(include="number")
    df_categorical = df.select_dtypes(include=["O"])
    numerical_stats = {
        k: stat(df_numerical).to_dict() for k, stat in NUMERICAL_STATS.items()
    }
    categorical_stats = {
        k: stat(df_categorical).to_dict() for k, stat in CATEGORICAL_STATS.items()
    }
    return {**numerical_stats, **categorical_stats}


def get_columnwise_comparison_stats(df_train: pd.DataFrame, df_generated: pd.DataFrame):
    df_numerical = df_train.select_dtypes(include="number")
    df_numerical_generated = df_generated.select_dtypes(include="number")

    numerical_stats = {
        k: {
            column: stat(df_numerical[column], df_numerical_generated[column])
            for column in df_train.columns
        }
        for k, stat in COLUMNWISE_NUMERICAL_COMPARISON_STATS.items()
    }

    return numerical_stats


def get_global_comparison_stats(df_train: pd.DataFrame, df_generated: pd.DataFrame):
    df_numerical = df_train.select_dtypes(include="number")
    df_numerical_generated = df_generated.select_dtypes(include="number")

    numerical_stats = {
        k: stat(df_numerical, df_numerical_generated)
        for k, stat in GLOBAL_NUMERICAL_COMPARISON_STATS.items()
    }

    return numerical_stats


def statistics_comparison(
    *,
    df_generated: pd.DataFrame,
    df_train: pd.DataFrame,
    path: Path,
    name_generated="generated",
    name_original="original",
    **kwargs,
):
    generated_stats = get_stats(df_generated)
    true_data_stats = get_stats(df_train)
    columnwise_comparison_stats = get_columnwise_comparison_stats(
        df_train=df_train, df_generated=df_generated
    )
    global_comparison_stats = (
        get_global_comparison_stats(df_train=df_train, df_generated=df_generated)
        if len(df_train.columns) > 1
        else None
    )
    stats = {
        name_generated: generated_stats,
        name_original: true_data_stats,
        "columnwise_comparison": columnwise_comparison_stats,
    }

    if global_comparison_stats:
        stats["global_comparison"] = global_comparison_stats
    store_json(stats, file=path / STATS_NAME)


def correlations_comparison(
    *,
    df_generated: pd.DataFrame,
    df_train: pd.DataFrame,
    path: Path,
    name_generated="generated",
    name_original="original",
    **kwargs,
):
    df_generated_numerical = df_generated.select_dtypes(include="number")
    df_train_numerical = df_train.select_dtypes(include="number")

    n = len(df_generated_numerical.columns)
    plt.figure().set_size_inches(n / 2 + 4, n / 2)

    if n < 2:
        return
    bi_corr = np.tril(df_generated_numerical.corr()) + np.triu(
        df_train_numerical.corr(), k=1
    )
    df = pd.DataFrame(
        bi_corr,
        index=df_generated_numerical.columns,
        columns=df_generated_numerical.columns,
    )
    sns.heatmap(
        df,
        annot=True,
        vmin=-1.0,
        vmax=1.0,
        cmap=sns.color_palette("coolwarm", as_cmap=True),
        alpha=0.8,
    )

    plt.title(f"Correlations (lower: {name_generated}, upper: {name_original})")
    plt.savefig(path / Path(f"{CORRELATIONS_NAME}.{IMAGE_FORMAT}"))
    plt.close()


def store_images(samples: torch.Tensor, *, folder: str) -> None:
    for i, sample in enumerate(samples):
        image.imsave(
            (f"{folder}/sample_{i}.{IMAGE_FORMAT}"),
            upscale(sample.numpy()),
            cmap="gray",
        )


def time_series_plot(
    x: np.ndarray, *, path: str, features_names: list, color: str
) -> None:
    k = len(features_names)
    aspect_ratio = 2
    h = 4
    fig, axes = plt.subplots(k, figsize=(h * aspect_ratio, h), sharex=True)
    for i, name in enumerate(features_names):
        axes[i].plot(x[:, :, i].T, alpha=0.5, color=color)
        axes[i].set_ylabel(name)
    plt.xlabel("t")
    fig.savefig(path)
    plt.close()


def time_series_plots(
    samples: np.ndarray,
    *,
    folder: str,
    n_plots: int,
    series_per_plot: int,
    features_names: list,
    color: str = "orange",
) -> None:
    assert samples.shape[-1] == len(features_names)

    for i in tqdm.tqdm(range(n_plots), desc="Generating time-series plots"):
        start = i * series_per_plot
        end = min(start + series_per_plot, len(samples))
        time_series_plot(
            x=samples[start:end],
            features_names=features_names,
            path=f"{folder}/{i}.{IMAGE_FORMAT}",
            color=color,
        )
    plt.close()


def score_plot(score_function, path: Path):
    plt.close()
    x = np.linspace(-2.0, 2.0, 200).reshape(-1, 1)
    x_tensor = torch.tensor(x, dtype=torch.float32, device=score_function.device)

    t_list = np.exp(np.linspace(np.log(MINIMUM_SAMPLING_T), 0, 20))
    for t in t_list:
        score = score_function(X=x_tensor, t=t).cpu()
        plt.plot(x[:], score[:], alpha=max((1 - t) ** 2, 0.05), color="blue")

    plt.plot(x, -x, label="N(0,1)", color="black")

    plt.plot(
        x,
        score_function(X=x_tensor, t=0.0).cpu(),
        alpha=1,
        color="red",
        label=f"t={0}",
        linewidth=2.0,
    )

    # x = x.squeeze()
    # plt.plot(x[:-1], mixture_score(x), label="mixture_score", alpha=0.8, color="green")
    plt.legend()
    plt.savefig(path / Path(f"score.{IMAGE_FORMAT}"))
    plt.close()


def satisfaction_plot(
    X: Tensor, constraint: Constraint, path: str | Path, label: str, **kwargs
) -> None:
    satisfaction = (torch.exp(constraint.f(X) * constraint.strength)).cpu().numpy()
    store_json(
        {
            "median": float(np.median(satisfaction)),
            "std": float(np.std(satisfaction)),
            "perfect_fraction": np.mean(satisfaction >= 0.5),
        },
        file=Path(path) / Path(f"satisfaction_{label}.json"),
    )
    plt.hist(satisfaction, alpha=0.8, **kwargs)
    plt.savefig(f"{str(path)}/constraint_satisfaction_{label}.{IMAGE_FORMAT}")
    plt.close()

src/report.py
DavideScassola/score-based-constrained-generation
from dataclasses import dataclass
from pathlib import Path

from src.util import load_module

from .data import Dataset
from .models.model import Model


@dataclass
class TrainConfig:
    dataset: Dataset
    model: Model
    generation_options: dict
    name: str | None = None


def getConfig(path: str | Path) -> TrainConfig:
    return load_module(str(path)).CONFIG

src/train_config.py
DavideScassola/score-based-constrained-generation
import os
import shutil
from datetime import datetime
from pathlib import Path

from src.constants import CONFIG_FILE_NAME, MODELS_FOLDER
from src.util import create_experiment_folder

from .train_config import TrainConfig, getConfig


def create_models_folder() -> Path:
    os.makedirs(MODELS_FOLDER, exist_ok=True)
    return MODELS_FOLDER


def create_model_folder() -> Path:
    models_folder = create_models_folder()
    folder_name = Path(datetime.now().strftime("%Y-%m-%d_%H:%M:%S_%f"))
    path = models_folder / folder_name
    os.makedirs(path)
    return path


def store_results(*, config_file: str | Path, config: TrainConfig):
    folder = Path(
        create_experiment_folder(
            path=create_models_folder(), postfix=getConfig(str(config_file)).name
        )
    )
    config.model.store(folder)
    shutil.copyfile(config_file, folder / CONFIG_FILE_NAME)
    config.model.generate_report(
        path=folder,
        dataset=config.dataset,
        generation_options=config.generation_options,
    )


def run_model_training(config_file: str):
    config = getConfig(config_file)
    config.model.train(config.dataset)
    store_results(config_file=config_file, config=config)

src/train_controller.py
DavideScassola/score-based-constrained-generation
import fnmatch
import importlib.util
import json
import math
import os
import random
from datetime import datetime
from pathlib import Path
from types import ModuleType
from typing import Callable, Tuple

import numpy as np
import torch


def load_json(file: str | Path) -> dict:
    with open(file) as json_file:
        d = json.load(json_file)
    return d


def store_json(d: dict, *, file: str | Path):
    with open(file, "w") as f:
        json.dump(d, f, indent=4)


def file_name(file: str | Path) -> str:
    return str(file).split("/")[-1]


def load_module(path: str | Path) -> ModuleType:
    spec = importlib.util.spec_from_file_location("module", path)
    if spec == None:
        raise ValueError(f"{path} is not a valid module path")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)  # type: ignore
    return module


def execute_python_script(path: str | Path):
    os.system(f"python {path}")


def create_experiment_folder(*, path: Path, postfix: str | None = None) -> Path:
    postfix = f"_{postfix}" if postfix else ""
    folder_name = Path(datetime.now().strftime("%Y-%m-%d_%H:%M:%S_%f") + postfix)
    experiment_folder = path / folder_name
    os.makedirs(experiment_folder)
    return experiment_folder


def gradient(
    *, f: Callable[[torch.Tensor], torch.Tensor], X: torch.Tensor
) -> torch.Tensor:
    # TODO: Check if all of this boilerplate is necessary
    was_grad_enabled = torch.is_grad_enabled()
    torch.set_grad_enabled(True)
    X.requires_grad = True
    X.grad = None

    value = f(X).sum()
    value.backward()
    assert torch.isfinite(value), f(X)
    grad = X.grad

    X.grad = None
    X.requires_grad = False
    torch.set_grad_enabled(was_grad_enabled)
    return grad  # type: ignore


def upscale(image, n=20):
    new_data = np.zeros(np.array(image.shape) * n)

    for j in range(image.shape[0]):
        for k in range(image.shape[1]):
            new_data[j * n : (j + 1) * n, k * n : (k + 1) * n] = image[j, k]

    return new_data


def find(folder: str, *, pattern: str = "*", index: int = -1):
    """
    Picks the last mnist stored model
    """
    matches = fnmatch.filter(sorted(os.listdir(folder)), pattern)
    return f"{folder}/{matches[index]}" if matches else None


def D_kl(x: np.ndarray, y: np.ndarray) -> float:
    bounds = (min(np.min(x), np.min(y)), max(np.max(x), np.max(y)))
    bins = min(len(x), len(y)) // 20

    def density_estimator(k):
        b = np.histogram(k, density=False, range=bounds, bins=bins)[0] + 1
        return b / np.sum(b)

    px = density_estimator(x)
    py = density_estimator(y)
    return np.sum(np.where(px != 0, px * (np.log(px) - np.log(py)), 0))


def max_symmetric_D_kl(x: np.ndarray, y: np.ndarray) -> float:
    bounds = (min(np.min(x), np.min(y)), max(np.max(x), np.max(y)))
    bins = min(len(x), len(y)) // 20

    def density_estimator(k):
        b = (
            np.histogram(k, density=False, range=bounds, bins=bins)[0] + 1
        )  # +1 in order to have numerical stability
        return b / np.sum(b)

    px = density_estimator(x)
    py = density_estimator(y)

    dkl = lambda p1, p2: np.sum(np.where(p1 != 0, p1 * (np.log(p1) - np.log(p2)), 0))
    return max(dkl(px, py), dkl(py, px))


def l1_divergence(x: np.ndarray, y: np.ndarray) -> float:
    bounds = (min(np.min(x), np.min(y)), max(np.max(x), np.max(y)))
    bins = int(np.sqrt(min(len(x), len(y))))

    def density_estimator(k):
        b = np.histogram(k, density=False, range=bounds, bins=bins)[0]
        return b / np.sum(b)

    px = density_estimator(x)
    py = density_estimator(y)

    return np.sum(np.abs(px - py)) / 2


def rejection_sampling(
    x: np.ndarray, acceptance_function, max_acceptance: float | None = 1.0
) -> np.ndarray:
    """acceptance_function is a probability density function"""
    acceptance = acceptance_function(x)
    observed_max_acceptance = np.max(acceptance)
    if not max_acceptance or observed_max_acceptance > max_acceptance:
        max_acceptance = observed_max_acceptance
    acceptance /= max_acceptance
    selection = np.random.binomial(n=1, p=acceptance) == 1
    return x[selection]


def reciprocal_distribution_sampler(
    *, low: float, up: float, size: tuple, dtype=None, device=None
) -> torch.Tensor:
    return torch.exp(
        torch.rand(size=size, dtype=dtype, device=device)
        * (math.log(up) - math.log(low))
        + math.log(low)
    )


def uniform_sampler(
    *, low: float, up: float, size: tuple, dtype=None, device=None
) -> torch.Tensor:
    return torch.rand(size=size, dtype=dtype, device=device) * (up - low) + low


def names_index_map(file) -> dict[str, int]:
    with open(file, "r") as f:
        names = (f.readline()[:-1]).split(",")
    return {n: i for (i, n) in enumerate(names)}


def set_seeds(seed: int):
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)

src/util.py
DavideScassola/score-based-constrained-generation