Update

app.py

@@ -1,16 +1,18 @@
 import numpy as np
 import open_clip
 import streamlit as st
-import torch
 
 from app_lib.main import main
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+if "sidebar_state" not in st.session_state:
+    st.session_state.sidebar_state = "collapsed"
+if "disabled" not in st.session_state:
+    st.session_state.disabled = False
+if "results" not in st.session_state:
+    st.session_state.results = None
+
+st.set_page_config(layout="wide", initial_sidebar_state=st.session_state.sidebar_state)
 
-st.set_page_config(
-    layout="wide",
-    initial_sidebar_state=st.session_state.get("sidebar_state", "collapsed"),
-)
 st.session_state.sidebar_state = "collapsed"
 st.markdown(
     """
@@ -21,6 +23,15 @@ st.markdown(
             input {
                 font-family: monospace !important;
             }
+
+            [data-testid="stHorizontalBlock"] {
+                align-items: center;
+            }
+            div.stSpinner > div {
+                text-align:center;
+                align-items: center;
+                justify-content: center;
+            }
         </style>
         """,
     unsafe_allow_html=True,
@@ -36,19 +47,5 @@ st.markdown(
         """,
 )
 
-
-def load_clip():
-    model, _, preprocess = open_clip.create_model_and_transforms(
-        "hf-hub:laion/CLIP-ViT-B-32-laion2B-s34B-b79K"
-    )
-    tokenizer = open_clip.get_tokenizer("hf-hub:laion/CLIP-ViT-B-32-laion2B-s34B-b79K")
-
-
-def test(
-    image, class_name, concepts, cardinality, model_name, dataset_name="imagenette"
-):
-    print("test!")
-
-
 if __name__ == "__main__":
     main()

app_lib/ckde.py

@@ -0,0 +1,77 @@
+import numpy as np
+import torch
+from scipy.spatial.distance import cdist
+from scipy.stats import gaussian_kde
+
+class cKDE:
+    def __init__(self, config, concept_class_name=None, concept_image_idx=None):
+        ckde_config = config.ckde
+        self.image_size = image_size = ckde_config.get("image_size", 128)
+        self.metric = ckde_config.get("metric", "euclidean")
+        self.scale_method = ckde_config.get("scale_method", "neff")
+        self.scale = ckde_config.get("scale", 2000)
+
+        self.Z = self.dataset.Z
+        self.H = self.dataset.H
+
+    def _quantile_scale(self, Z_cond_dist):
+        return np.quantile(Z_cond_dist, self.scale)
+
+    def _neff_scale(self, Z_cond_dist):
+        scales = np.linspace(1e-02, 0.4, 100)[:, None]
+
+        _Z_cond_dist = np.tile(Z_cond_dist, (len(scales), 1))
+
+        weights = np.exp(-(_Z_cond_dist**2) / (2 * scales**2))
+        neff = (np.sum(weights, axis=1) ** 2) / np.sum(weights**2, axis=1)
+        diff = np.abs(neff - self.scale)
+        scale_idx = np.argmin(diff)
+        return scales[scale_idx].item()
+
+    def _sample(self, z, cond_idx, m):
+        sample_idx = list(set(range(len(z))) - set(cond_idx))
+
+        kde, _ = self.kde(z, cond_idx)
+
+        sample_z = np.tile(z, (m, 1))
+        sample_z[:, sample_idx] = kde.resample(m).T
+
+        return sample_z
+
+    def kde(self, z, cond_idx):
+        sample_idx = list(set(range(len(z))) - set(cond_idx))
+
+        Z_sample = self.Z[:, sample_idx]
+        Z_cond = self.Z[:, cond_idx]
+
+        z_cond = z[cond_idx]
+        Z_cond_dist = cdist(z_cond.reshape(1, -1), Z_cond, self.metric).squeeze()
+
+        if self.scale_method == "constant":
+            scale = self.scale
+        if self.scale_method == "quantile":
+            scale = self._quantile_scale(Z_cond_dist)
+        elif self.scale_method == "neff":
+            scale = self._neff_scale(Z_cond_dist)
+
+        weights = np.exp(-(Z_cond_dist**2) / (2 * scale**2))
+
+        return gaussian_kde(Z_sample.T, weights=weights), scale
+
+    def nearest_neighbor(self, z):
+        dist = cdist(z, self.Z, metric=self.metric)
+        return np.argmin(dist, axis=-1)
+
+    def sample(self, z, cond_idx, m=1, return_images=False):
+        if z.ndim == 1:
+            z = z.reshape(1, -1)
+
+        sample_z = np.concatenate([self._sample(_z, cond_idx, m) for _z in z], axis=0)
+
+        nn_idx = self.nearest_neighbor(sample_z)
+        sample_h = self.H[nn_idx]
+
+        if return_images:
+            sample_images = torch.stack([self.dataset[_idx][0] for _idx in nn_idx])
+            return sample_z, sample_h, sample_images
+        return sample_z, sample_h

app_lib/main.py

@@ -1,4 +1,6 @@
+import torch
 import streamlit as st
+import time
 
 from app_lib.user_input import (
     get_cardinality,
@@ -7,9 +9,20 @@ from app_lib.user_input import (
     get_image,
     get_model_name,
 )
+from app_lib.test import (
+    load_dataset,
+    load_model,
+    encode_image,
+    encode_concepts,
+    encode_class_name,
+)
+
+
+def _disable():
+    st.session_state.disabled = True
 
 
-def main():
+def main(device=torch.device("cuda" if torch.cuda.is_available() else "cpu")):
     columns = st.columns([0.40, 0.60])
 
     with columns[0]:
@@ -27,7 +40,11 @@ def main():
             cardinality = get_cardinality(concepts, concepts_ready)
 
         with row2[0]:
-            change_image_button = st.button("Change Image", use_container_width=True)
+            change_image_button = st.button(
+                "Change Image",
+                use_container_width=True,
+                disabled=st.session_state.disabled,
+            )
             if change_image_button:
                 st.session_state.sidebar_state = "expanded"
                 st.experimental_rerun()
@@ -39,13 +56,38 @@ def main():
                 error_message += f"- {class_error}\n"
             if concepts_error is not None:
                 error_message += f"- {concepts_error}\n"
+            if error_message:
+                st.error(error_message)
 
             test_button = st.button(
                 "Test",
-                help=None if ready else error_message,
                 use_container_width=True,
-                disabled=not ready,
+                on_click=_disable,
+                disabled=st.session_state.disabled or not ready,
             )
 
-            if test_button:
-                test(image, class_name, concepts, cardinality, model_name)
+    with columns[1]:
+        if test_button:
+            with st.spinner("Loading dataset"):
+                embedding = load_dataset("imagenette", model_name)
+                time.sleep(1)
+
+            with st.spinner("Loading model"):
+                model, preprocess, tokenizer = load_model(model_name, device)
+                time.sleep(1)
+
+            with st.spinner("Encoding concepts"):
+                cbm = encode_concepts(tokenizer, model, concepts, device)
+                time.sleep(1)
+
+            with st.spinner("Preparing zero-shot classifier"):
+                classifier = encode_class_name(tokenizer, model, class_name, device)
+
+            with st.spinner("Encoding image"):
+                h = encode_image(model, preprocess, image, device)
+                z = h @ cbm.T
+                print(h.shape, cbm.shape, z.shape)
+                time.sleep(2)
+
+            st.session_state.disabled = False
+            st.experimental_rerun()

app_lib/test.py

@@ -0,0 +1,84 @@
+import torch
+import clip
+import open_clip
+import h5py
+
+from huggingface_hub import hf_hub_download
+
+from app_lib.utils import SUPPORTED_MODELS
+
+
+def _get_open_clip_model(model_name, device):
+    backbone = model_name.split(":")[-1]
+
+    model, _, preprocess = open_clip.create_model_and_transforms(
+        SUPPORTED_MODELS[model_name], device=device
+    )
+    model.eval()
+    tokenizer = open_clip.get_tokenizer(backbone)
+    return model, preprocess, tokenizer
+
+
+def _get_clip_model(model_name, device):
+    backbone = model_name.split(":")[-1]
+    model, preprocess = clip.load(backbone, device=device)
+    tokenizer = clip.tokenize
+    return model, preprocess, tokenizer
+
+
+def load_dataset(dataset_name, model_name):
+    dataset_path = hf_hub_download(
+        repo_id="jacopoteneggi/IBYDMT",
+        filename=f"{dataset_name}_{model_name}_train.h5",
+        repo_type="dataset",
+    )
+
+    with h5py.File(dataset_path, "r") as dataset:
+        embedding = dataset["embedding"][:]
+    return embedding
+
+
+def load_model(model_name, device):
+    print(model_name)
+    if "open_clip" in model_name:
+        model, preprocess, tokenizer = _get_open_clip_model(model_name, device)
+    elif "clip" in model_name:
+        model, preprocess, tokenizer = _get_clip_model(model_name, device)
+    return model, preprocess, tokenizer
+
+
+@torch.no_grad()
+@torch.cuda.amp.autocast()
+def encode_concepts(tokenizer, model, concepts, device):
+    concepts_text = tokenizer(concepts).to(device)
+
+    concept_features = model.encode_text(concepts_text)
+    concept_features /= torch.linalg.norm(concept_features, dim=-1, keepdim=True)
+    return concept_features.cpu().numpy()
+
+
+@torch.no_grad()
+@torch.cuda.amp.autocast()
+def encode_image(model, preprocess, image, device):
+    image = preprocess(image)
+    image = image.unsqueeze(0)
+    image = image.to(device)
+
+    image_features = model.encode_image(image)
+    image_features /= image_features.norm(dim=-1, keepdim=True)
+    return image_features.cpu().numpy()
+
+
+@torch.no_grad()
+@torch.cuda.amp.autocast()
+def encode_class_name(tokenizer, model, class_name, device):
+    class_text = tokenizer([f"A photo of a {class_name}"]).to(device)
+
+    class_features = model.encode_text(class_text)
+    class_features /= torch.linalg.norm(class_features, dim=-1, keepdim=True)
+    return class_features.cpu().numpy()
+
+
+def test(image, class_name, concepts, cardinality, dataset_name, model_name, device):
+    model, preprocess = load_model(model_name, device)
+    print(f"loaded {model_name}")

app_lib/user_input.py

@@ -24,8 +24,9 @@ def _validate_concepts(concepts):
 def get_model_name():
     return st.selectbox(
         "Choose a model to test",
-        options=SUPPORTED_MODELS,
+        options=list(SUPPORTED_MODELS.keys()),
         help="Name of the vision-language model to test the predictions of.",
+        disabled=st.session_state.disabled,
     )
 
 
@@ -49,6 +50,7 @@ def get_class_name():
         "Class to test",
         help="Name of the class to build the zero-shot CLIP classifier with.",
         value="cat",
+        disabled=st.session_state.disabled,
     )
 
     class_ready, class_error = _validate_class_name(class_name)
@@ -61,6 +63,7 @@ def get_concepts():
         help="List of concepts to test the predictions of the model with. Write one concept per line.",
         height=160,
         value="piano\ncute\nwhiskers\nmusic\nwild",
+        disabled=st.session_state.disabled,
     )
     concepts = concepts.split("\n")
     concepts = [concept.strip() for concept in concepts]
@@ -79,5 +82,5 @@ def get_cardinality(concepts, concepts_ready):
         max_value=max(2, len(concepts) - 1),
         value=1,
         step=1,
-        disabled=not concepts_ready,
+        disabled=st.session_state.disabled or not concepts_ready,
     )

app_lib/utils.py

@@ -5,10 +5,22 @@ supported_models_path = hf_hub_download(
     filename="supported_models.txt",
     repo_type="dataset",
 )
+supported_datasets_path = hf_hub_download(
+    repo_id="jacopoteneggi/IBYDMT",
+    filename="supported_datasets.txt",
+    repo_type="dataset",
+)
 
-SUPPORTED_MODELS = []
+SUPPORTED_MODELS = {}
 with open(supported_models_path, "r") as f:
     for line in f:
         line = line.strip()
-        model_name, _ = line.split(",")
-        SUPPORTED_MODELS.append(model_name)
+        model_name, model_url = line.split(",")
+        SUPPORTED_MODELS[model_name] = model_url
+
+
+SUPPORTED_DATASETS = []
+with open(supported_models_path, "r") as f:
+    for line in f:
+        dataset_name = line.strip()
+        SUPPORTED_DATASETS.append(dataset_name)

ibydmt/__init__.py

@@ -0,0 +1 @@
+from ibydmt.test import SKIT, cSKIT, xSKIT

ibydmt/bet.py

@@ -0,0 +1,56 @@
+from abc import ABC, abstractmethod
+from typing import Dict
+
+import numpy as np
+
+from ibydmt.utils import _get_cls, _register_cls
+
+
+class Bet(ABC):
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def compute(self, *args, **kwargs):
+        pass
+
+
+_BETS: Dict[str, Bet] = {}
+
+
+def register_bet(name):
+    return _register_cls(name, dict=_BETS)
+
+
+def get_bet(name):
+    return _get_cls(name, dict=_BETS)
+
+
+@register_bet("sign")
+class Sign(Bet):
+    def __init__(self, config):
+        super().__init__()
+        self.m = config.get("m", 0.5)
+        self.prev_g = []
+
+    def compute(self, g):
+        return self.m * np.sign(g)
+
+
+@register_bet("tanh")
+class Tanh(Bet):
+    def __init__(self, config):
+        super().__init__()
+        self.alpha = config.get("alpha", 0.20)
+        self.prev_g = []
+
+    def compute(self, g):
+        if len(self.prev_g) < 2:
+            scale = 1
+        else:
+            l, u = np.quantile(self.prev_g, [self.alpha / 2, 1 - self.alpha / 2])
+            scale = u - l
+
+        self.prev_g.append(g)
+
+        return np.tanh(g / np.clip(scale, 1e-04, None))

ibydmt/payoff.py

@@ -0,0 +1,160 @@
+from abc import ABC, abstractmethod
+from functools import reduce
+
+import numpy as np
+from sklearn.metrics import pairwise_distances
+from sklearn.metrics.pairwise import linear_kernel, rbf_kernel
+
+from ibydmt.bet import get_bet
+
+
+class Payoff(ABC):
+    def __init__(self, config):
+        self.bet = get_bet(config.bet)(config)
+
+    @abstractmethod
+    def compute(self, *args, **kwargs):
+        pass
+
+
+class Kernel:
+    def __init__(self, kernel: str, scale_method: str, scale: float):
+        if kernel == "linear":
+            self.base_kernel = linear_kernel
+        elif kernel == "rbf":
+            self.base_kernel = rbf_kernel
+
+            self.scale_method = scale_method
+            self.scale = scale
+
+            self.gamma = None
+            self.recompute_gamma = True
+            self.prev = None
+        else:
+            raise NotImplementedError(f"{kernel} is not implemented")
+
+    def __call__(self, x, y):
+        if self.base_kernel == linear_kernel:
+            return self.base_kernel(x, y)
+        if self.base_kernel == rbf_kernel:
+            if self.scale_method == "constant":
+                self.gamma = self.scale
+            elif self.scale_method == "quantile":
+                if self.prev is None:
+                    self.prev = y
+
+                if self.recompute_gamma:
+                    dist = pairwise_distances(
+                        self.prev.reshape(-1, self.prev.shape[-1])
+                    )
+                    scale = np.quantile(dist, self.scale)
+                    gamma = 1 / (2 * scale**2) if scale > 0 else None
+                    self.gamma = gamma
+
+                    if len(self.prev) > 100:
+                        self.recompute_gamma = False
+                    self.prev = np.vstack([self.prev, x])
+            else:
+                raise NotImplementedError(
+                    f"{self.scale} is not implemented for rbf_kernel"
+                )
+            return self.base_kernel(x, y, gamma=self.gamma)
+
+
+class KernelPayoff(Payoff):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.kernel = config.kernel
+        self.scale_method = config.get("kernel_scale_method", "quantile")
+        self.scale = config.get("kernel_scale", 0.5)
+
+    @abstractmethod
+    def witness_function(self, d, prev_d):
+        pass
+
+    def compute(self, d, null_d, prev_d):
+        g = reduce(
+            lambda acc, u: acc
+            + self.witness_function(u[0], prev_d)
+            - self.witness_function(u[1], prev_d),
+            zip(d, null_d),
+            0,
+        )
+
+        return self.bet.compute(g)
+
+
+class HSIC(KernelPayoff):
+    def __init__(self, config):
+        super().__init__(config)
+
+        kernel = self.kernel
+        scale_method = self.scale_method
+        scale = self.scale
+
+        self.kernel_y = Kernel(kernel, scale_method, scale)
+        self.kernel_z = Kernel(kernel, scale_method, scale)
+
+    def witness_function(self, d, prev_d):
+        y, z = d
+        prev_y, prev_z = prev_d[:, 0], prev_d[:, 1]
+
+        y_mat = self.kernel_y(y.reshape(-1, 1), prev_y.reshape(-1, 1))
+        z_mat = self.kernel_z(z.reshape(-1, 1), prev_z.reshape(-1, 1))
+
+        mu_joint = np.mean(y_mat * z_mat)
+        mu_prod = np.mean(y_mat, axis=1) @ np.mean(z_mat, axis=1)
+        return mu_joint - mu_prod
+
+
+class cMMD(KernelPayoff):
+    def __init__(self, config):
+        super().__init__(config)
+
+        kernel = self.kernel
+        scale_method = self.scale_method
+        scale = self.scale
+
+        self.kernel_y = Kernel(kernel, scale_method, scale)
+        self.kernel_zj = Kernel(kernel, scale_method, scale)
+        self.kernel_cond_z = Kernel(kernel, scale_method, scale)
+
+    def witness_function(self, u, prev_d):
+        y, zj, cond_z = u[0], u[1], u[2:]
+
+        prev_y, prev_zj, prev_null_zj, prev_cond_z = (
+            prev_d[:, 0],
+            prev_d[:, 1],
+            prev_d[:, 2],
+            prev_d[:, 3:],
+        )
+
+        y_mat = self.kernel_y(y.reshape(-1, 1), prev_y.reshape(-1, 1))
+        zj_mat = self.kernel_zj(zj.reshape(-1, 1), prev_zj.reshape(-1, 1))
+        cond_z_mat = self.kernel_cond_z(
+            cond_z.reshape(-1, prev_cond_z.shape[1]),
+            prev_cond_z.reshape(-1, prev_cond_z.shape[1]),
+        )
+
+        null_zj_mat = self.kernel_zj(zj.reshape(-1, 1), prev_null_zj.reshape(-1, 1))
+
+        mu = np.mean(y_mat * zj_mat * cond_z_mat)
+        mu_null = np.mean(y_mat * null_zj_mat * cond_z_mat)
+        return mu - mu_null
+
+
+class xMMD(KernelPayoff):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.kernel = Kernel(self.kernel, self.scale_method, self.scale)
+
+    def witness_function(self, u, prev_d):
+        prev_y, prev_y_null = prev_d[:, 0], prev_d[:, 1]
+
+        mu_y = np.mean(self.kernel(u.reshape(-1, 1), prev_y.reshape(-1, 1)), axis=1)
+        mu_y_null = np.mean(
+            self.kernel(u.reshape(-1, 1), prev_y_null.reshape(-1, 1)), axis=1
+        )
+        return mu_y - mu_y_null

ibydmt/test.py

@@ -0,0 +1,159 @@
+import functools
+from abc import ABC, abstractmethod
+from collections import deque
+from typing import Callable, Tuple, Union
+
+import numpy as np
+import torch
+from jaxtyping import Float
+
+from ibydmt.payoff import HSIC, cMMD, xMMD
+from ibydmt.wealth import get_wealth
+
+Array = Union[np.ndarray, torch.Tensor]
+
+
+class Tester(ABC):
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def test(self, *args, **kwargs) -> Tuple[bool, int]:
+        pass
+
+
+class SequentialTester(Tester):
+    def __init__(self, config):
+        super().__init__()
+        self.wealth = get_wealth(config.wealth)(config)
+
+        self.tau_max = config.tau_max
+
+
+class SKIT(SequentialTester):
+    """Global Independence Tester"""
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.payoff = HSIC(config)
+
+    def test(self, Y: Float[Array, "N"], Z: Float[Array, "N"]) -> Tuple[bool, int]:
+        D = np.stack([Y, Z], axis=1)
+        for t in range(1, self.tau_max):
+            d = D[2 * t : 2 * (t + 1)]
+            prev_d = D[: 2 * t]
+
+            null_d = np.stack([d[:, 0], np.flip(d[:, 1])], axis=1)
+
+            payoff = self.payoff.compute(d, null_d, prev_d)
+            self.wealth.update(payoff)
+
+            if self.wealth.rejected:
+                return (True, t)
+        return (False, t)
+
+
+class cSKIT(SequentialTester):
+    """Global Conditional Independence Tester"""
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.payoff = cMMD(config)
+
+    def _sample(
+        self,
+        z: Float[Array, "N D"],
+        j: int = None,
+        cond_p: Callable[[Float[Array, "N D"], list[int]], Float[Array, "N D"]] = None,
+    ) -> Tuple[Float[Array, "N"], Float[Array, "N"], Float[Array, "N D-1"]]:
+        C = list(set(range(z.shape[1])) - {j})
+
+        zj, cond_z = z[:, [j]], z[:, C]
+        samples = cond_p(z, C)
+        null_zj = samples[:, [j]]
+        return zj, null_zj, cond_z
+
+    def test(
+        self,
+        Y: Float[Array, "N"],
+        Z: Float[Array, "N D"],
+        j: int,
+        cond_p: Callable[[Float[Array, "N D"], list[int]], Float[Array, "N D"]],
+    ) -> Tuple[bool, int]:
+        sample = functools.partial(self._sample, j=j, cond_p=cond_p)
+
+        prev_y, prev_z = Y[:1][:, None], Z[:1]
+        prev_zj, prev_null_zj, prev_cond_z = sample(prev_z)
+        prev_d = np.concatenate([prev_y, prev_zj, prev_null_zj, prev_cond_z], axis=-1)
+        for t in range(1, self.tau_max):
+            y, z = Y[[t]][:, None], Z[[t]]
+            zj, null_zj, cond_z = sample(z)
+
+            u = np.concatenate([y, zj, cond_z], axis=-1)
+            null_u = np.concatenate([y, null_zj, cond_z], axis=-1)
+
+            payoff = self.payoff.compute(u, null_u, prev_d)
+            self.wealth.update(payoff)
+
+            d = np.concatenate([y, zj, null_zj, cond_z], axis=-1)
+            prev_d = np.vstack([prev_d, d])
+
+            if self.wealth.rejected:
+                return (True, t)
+        return (False, t)
+
+
+class xSKIT(SequentialTester):
+    """Local Conditional Independence Tester"""
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.payoff = xMMD(config)
+
+        self._queue = deque()
+
+    def _sample(
+        self,
+        z: Float[Array, "D"],
+        j: int,
+        C: list[int],
+        cond_p: Callable[[Float[Array, "D"], list[int], int], Float[Array, "N D2"]],
+        model: Callable[[Float[Array, "N D2"]], Float[Array, "N"]],
+    ) -> Tuple[Float[Array, "1"], Float[Array, "1"]]:
+
+        if len(self._queue) == 0:
+            Cuj = C + [j]
+
+            h = cond_p(z, Cuj, self.tau_max)
+            null_h = cond_p(z, C, self.tau_max)
+
+            y = model(h)[:, None]
+            null_y = model(null_h)[:, None]
+
+            self._queue.extend(zip(y, null_y))
+
+        return self._queue.pop()
+
+    def test(
+        self,
+        z: Float[Array, "D"],
+        j: int,
+        C: list[int],
+        cond_p: Callable[[Float[Array, "D"], list[int], int], Float[Array, "N D2"]],
+        model: Callable[[Float[Array, "N D2"]], Float[Array, "N"]],
+    ) -> Tuple[bool, int]:
+        sample = functools.partial(self._sample, z, j, C, cond_p, model)
+
+        prev_d = np.stack(sample(), axis=1)
+        for t in range(1, self.tau_max):
+            y, null_y = sample()
+
+            payoff = self.payoff.compute(y, null_y, prev_d)
+            self.wealth.update(payoff)
+
+            d = np.stack([y, null_y], axis=1)
+            prev_d = np.vstack([prev_d, d])
+
+            if self.wealth.rejected:
+                return (True, t)
+        return (False, t)

ibydmt/utils.py

@@ -0,0 +1,12 @@
+def _register_cls(name, dict=None):
+    def _register(cls):
+        if name in dict:
+            raise ValueError(f"{name} is already registered")
+
+        dict[name] = cls
+
+    return _register
+
+
+def _get_cls(name, dict=None):
+    return dict[name]

ibydmt/wealth.py

@@ -0,0 +1,72 @@
+from abc import ABC, abstractmethod
+from typing import Dict
+
+import numpy as np
+
+from ibydmt.utils import _get_cls, _register_cls
+
+
+class Wealth(ABC):
+    def __init__(self, config):
+        self.significance_level = config.significance_level
+        self.rejected = False
+
+    @abstractmethod
+    def update(self, payoff):
+        pass
+
+
+_WEALTH: Dict[str, Wealth] = {}
+
+
+def register_wealth(name):
+    return _register_cls(name, dict=_WEALTH)
+
+
+def get_wealth(name):
+    return _get_cls(name, dict=_WEALTH)
+
+
+@register_wealth("mixture")
+class Mixture(Wealth):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.grid_size = grid_size = config.grid_size
+        self.wealth = np.ones((grid_size,))
+        self.wealth_flag = np.ones(grid_size, dtype=bool)
+        self.v = np.linspace(0.05, 0.95, grid_size)
+
+    def update(self, payoff):
+        raise NotImplementedError
+
+
+@register_wealth("ons")
+class ONS(Wealth):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.w = 1.0
+        self.v = 0
+        self.a = 1
+
+        self.min_v, self.max_v = config.get("min_v", 0), config.get("max_v", 1 / 2)
+        self.wealth_flag = False
+
+    def _update_v(self, payoff):
+        z = payoff / (1 + self.v * payoff)
+        self.a += z**2
+        self.v = max(
+            self.min_v, min(self.max_v, self.v + 2 / (2 - np.log(3)) * z / self.a)
+        )
+
+    def update(self, payoff):
+        w = self.w * (1 + self.v * payoff)
+
+        if w >= 0 and not self.wealth_flag:
+            self.w = w
+            if self.w >= 1 / self.significance_level:
+                self.rejected = True
+            self._update_v(payoff)
+        else:
+            self.wealth_flag = True

requirements.txt

@@ -1,3 +1,3 @@
 streamlit-image-select
 clip @ git+https://github.com/openai/CLIP@main
-open_clip_torch
+open_clip_torch