gensbi-examples 0.0.2__py3-none-any.whl

gensbi_examples/tasks.py

@@ -0,0 +1,343 @@
+import jax
+from jax import numpy as jnp
+import grain
+import numpy as np
+
+from datasets import load_dataset
+from huggingface_hub import hf_hub_download
+import json
+
+# from .utils import download_artifacts
+from .graph import faithfull_mask, min_faithfull_mask, moralize
+
+
+def process_joint(batch):
+    cond = batch["xs"][..., None]
+    obs = batch["thetas"][..., None]
+    data = np.concatenate((obs, cond), axis=1)
+    return data
+
+
+def process_conditional(batch):
+    cond = batch["xs"][..., None]
+    obs = batch["thetas"][..., None]
+    return obs, cond
+
+
+class Task:
+    def __init__(self, task_name, kind="joint"):
+
+        self.repo_name = "aurelio-amerio/SBI-benchmarks"
+
+        self.task_name = task_name
+
+        fname = hf_hub_download(
+            repo_id=self.repo_name, filename="metadata.json", repo_type="dataset"
+        )
+        with open(fname, "r") as f:
+            metadata = json.load(f)
+
+        self.dataset = load_dataset(self.repo_name, task_name).with_format("numpy")
+        self.dataset_posterior = load_dataset(
+            self.repo_name, f"{task_name}_posterior"
+        ).with_format("numpy")
+
+        self.max_samples = self.dataset["train"].num_rows
+
+        self.observations = self.dataset_posterior["reference_posterior"][
+            "observations"
+        ]
+        self.reference_samples = self.dataset_posterior["reference_posterior"][
+            "reference_samples"
+        ]
+
+        self.true_parameters = self.dataset_posterior["reference_posterior"][
+            "true_parameters"
+        ]
+
+        self.dim_cond = metadata[task_name]["dim_cond"]
+        self.dim_obs = metadata[task_name]["dim_obs"]
+
+        self.dim_joint = self.dim_cond + self.dim_obs
+
+        self.num_observations = len(self.observations)
+        self.kind = kind
+
+        if kind == "joint":
+            self.process_fn = process_joint
+        elif kind == "conditional":
+            self.process_fn = process_conditional
+        else:
+            raise ValueError(f"Unknown kind: {kind}")
+
+    def get_train_dataset(self, batch_size, nsamples=1e5):
+        assert (
+            nsamples < self.max_samples
+        ), f"nsamples must be less than {self.max_samples}"
+
+        df = self.dataset["train"].select(range(int(nsamples)))  # [:]
+
+        dataset_grain = (
+            grain.MapDataset.source(df).shuffle(42).repeat().to_iter_dataset()
+        )
+
+        performance_config = grain.experimental.pick_performance_config(
+            ds=dataset_grain,
+            ram_budget_mb=1024 * 4,
+            max_workers=None,
+            max_buffer_size=None,
+        )
+
+        dataset_batched = (
+            dataset_grain.batch(batch_size)
+            .map(self.process_fn)
+            .mp_prefetch(performance_config.multiprocessing_options)
+        )
+
+        return dataset_batched
+
+    def get_val_dataset(self, batch_size):
+        df = self.dataset["validation"]  # [:]
+
+        val_dataset_grain = (
+            grain.MapDataset.source(df).shuffle(42).repeat().to_iter_dataset()
+        )
+        performance_config = grain.experimental.pick_performance_config(
+            ds=val_dataset_grain,
+            ram_budget_mb=1024 * 4,
+            max_workers=None,
+            max_buffer_size=None,
+        )
+        val_dataset_grain = (
+            val_dataset_grain.batch(batch_size)
+            .map(self.process_fn)
+            .mp_prefetch(performance_config.multiprocessing_options)
+        )
+
+        return val_dataset_grain
+
+    def get_test_dataset(self, batch_size):
+        df = self.dataset["test"]  # [:]
+
+        val_dataset_grain = (
+            grain.MapDataset.source(df)
+            .shuffle(42)
+            .repeat()
+            .to_iter_dataset()
+            .batch(batch_size)
+            .map(self.process_fn)
+        )
+
+        return val_dataset_grain
+
+    def get_reference(self, num_observation=1):
+        """
+        Returns the reference posterior samples for a given number of observations.
+        """
+        if num_observation < 1 or num_observation > self.num_observations:
+            raise ValueError(
+                f"num_observation must be between 1 and {self.num_observations}"
+            )
+        obs = self.observations[num_observation - 1]
+        samples = self.reference_samples[num_observation - 1]
+        return obs, samples
+
+    def get_true_parameters(self, num_observation=1):
+        """
+        Returns the true parameters for a given number of observations.
+        """
+        if num_observation < 1 or num_observation > self.num_observations:
+            raise ValueError(
+                f"num_observation must be between 1 and {self.num_observations}"
+            )
+        return self.true_parameters[num_observation - 1]
+
+    def get_base_mask_fn(self):
+        raise NotImplementedError()
+
+    def get_edge_mask_fn(self, name="undirected"):
+        if name.lower() == "faithfull":
+            base_mask_fn = self.get_base_mask_fn()
+
+            def faithfull_edge_mask(node_id, condition_mask, meta_data=None):
+                base_mask = base_mask_fn(node_id, meta_data)
+                return faithfull_mask(base_mask, condition_mask)
+
+            return faithfull_edge_mask
+        elif name.lower() == "min_faithfull":
+            base_mask_fn = self.get_base_mask_fn()
+
+            def min_faithfull_edge_mask(node_id, condition_mask, meta_data=None):
+                base_mask = base_mask_fn(node_id, meta_data)
+
+                return min_faithfull_mask(base_mask, condition_mask)
+
+            return min_faithfull_edge_mask
+        elif name.lower() == "undirected":
+            base_mask_fn = self.get_base_mask_fn()
+
+            def undirected_edge_mask(node_id, condition_mask, meta_data=None):
+                base_mask = base_mask_fn(node_id, meta_data)
+                return moralize(base_mask)
+
+            return undirected_edge_mask
+
+        elif name.lower() == "directed":
+            base_mask_fn = self.get_base_mask_fn()
+
+            def directed_edge_mask(node_id, condition_mask, meta_data=None):
+                base_mask = base_mask_fn(node_id, meta_data)
+                return base_mask
+
+            return directed_edge_mask
+        elif name.lower() == "none":
+            return lambda node_id, condition_mask, *args, **kwargs: None
+        else:
+            raise NotImplementedError()
+
+
+class TwoMoons(Task):
+    def __init__(self, kind="joint"):
+        task_name = "two_moons"
+        super().__init__(task_name, kind=kind)
+
+    def get_base_mask_fn(self):
+        theta_dim = self.dim_obs
+        x_dim = self.dim_cond
+        thetas_mask = jnp.eye(theta_dim, dtype=jnp.bool_)
+        x_mask = jnp.tril(jnp.ones((theta_dim, x_dim), dtype=jnp.bool_))
+        base_mask = jnp.block(
+            [
+                [thetas_mask, jnp.zeros((theta_dim, x_dim))],
+                [jnp.ones((x_dim, theta_dim)), x_mask],
+            ]
+        )
+        base_mask = base_mask.astype(jnp.bool_)
+
+        def base_mask_fn(node_ids, node_meta_data):
+            return base_mask[node_ids, :][:, node_ids]
+
+        return base_mask_fn
+
+
+class BernoulliGLM(Task):
+    def __init__(self, kind="joint"):
+        task_name = "bernoulli_glm"
+        super().__init__(task_name, kind=kind)
+
+    def get_base_mask_fn(self):
+        raise NotImplementedError()
+
+
+class GaussianLinear(Task):
+    def __init__(self, kind="joint"):
+        task_name = "gaussian_linear"
+        super().__init__(task_name, kind=kind)
+
+    def get_base_mask_fn(self):
+        theta_dim = self.dim_obs
+        x_dim = self.dim_cond
+        thetas_mask = jnp.eye(theta_dim, dtype=jnp.bool_)
+        x_i_mask = jnp.eye(x_dim, dtype=jnp.bool_)
+        base_mask = jnp.block(
+            [[thetas_mask, jnp.zeros((theta_dim, x_dim))], [jnp.eye((x_dim)), x_i_mask]]
+        )
+        base_mask = base_mask.astype(jnp.bool_)
+
+        def base_mask_fn(node_ids, node_meta_data):
+            return base_mask[node_ids, :][:, node_ids]
+
+        return base_mask_fn
+
+
+class GaussianLinearUniform(Task):
+    def __init__(self, kind="joint"):
+        task_name = "gaussian_linear_uniform"
+        super().__init__(task_name, kind=kind)
+
+    def get_base_mask_fn(self):
+        theta_dim = self.dim_obs
+        x_dim = self.dim_cond
+        thetas_mask = jnp.eye(theta_dim, dtype=jnp.bool_)
+        x_i_mask = jnp.eye(x_dim, dtype=jnp.bool_)
+        base_mask = jnp.block(
+            [[thetas_mask, jnp.zeros((theta_dim, x_dim))], [jnp.eye((x_dim)), x_i_mask]]
+        )
+        base_mask = base_mask.astype(jnp.bool_)
+
+        def base_mask_fn(node_ids, node_meta_data):
+            return base_mask[node_ids, :][:, node_ids]
+
+        return base_mask_fn
+
+
+class GaussianMixture(Task):
+    def __init__(self, kind="joint"):
+        task_name = "gaussian_mixture"
+        super().__init__(task_name, kind=kind)
+
+    def get_base_mask_fn(self):
+        theta_dim = self.dim_obs
+        x_dim = self.dim_cond
+        thetas_mask = jnp.eye(theta_dim, dtype=jnp.bool_)
+        x_mask = jnp.tril(jnp.ones((theta_dim, x_dim), dtype=jnp.bool_))
+        base_mask = jnp.block(
+            [
+                [thetas_mask, jnp.zeros((theta_dim, x_dim))],
+                [jnp.ones((x_dim, theta_dim)), x_mask],
+            ]
+        )
+        base_mask = base_mask.astype(jnp.bool_)
+
+        def base_mask_fn(node_ids, node_meta_data):
+            return base_mask[node_ids, :][:, node_ids]
+
+        return base_mask_fn
+
+
+class SLCP(Task):
+    def __init__(self, kind="joint"):
+        task_name = "slcp"
+        super().__init__(task_name, kind=kind)
+
+    def get_base_mask_fn(self):
+        theta_dim = self.dim_obs
+        x_dim = self.dim_cond
+        thetas_mask = jnp.eye(theta_dim, dtype=jnp.bool_)
+        x_i_dim = x_dim // 4
+        x_i_mask = jax.scipy.linalg.block_diag(
+            *tuple([jnp.tril(jnp.ones((x_i_dim, x_i_dim), dtype=jnp.bool_))] * 4)
+        )
+        base_mask = jnp.block(
+            [
+                [thetas_mask, jnp.zeros((theta_dim, x_dim))],
+                [jnp.ones((x_dim, theta_dim)), x_i_mask],
+            ]
+        )
+        base_mask = base_mask.astype(jnp.bool_)
+
+        def base_mask_fn(node_ids, node_meta_data):
+            return base_mask[node_ids, :][:, node_ids]
+
+        return base_mask_fn
+
+
+def get_task(task_name, kind="joint"):
+    """
+    Returns a Task object based on the task name.
+    """
+    task_name = task_name.lower()
+    if task_name == "two_moons":
+        return TwoMoons(kind=kind)
+    elif task_name == "bernoulli_glm":
+        return BernoulliGLM(kind=kind)
+    elif task_name == "gaussian_linear":
+        return GaussianLinear(kind=kind)
+    elif task_name == "gaussian_linear_uniform":
+        return GaussianLinearUniform(kind=kind)
+    elif task_name == "gaussian_mixture":
+        return GaussianMixture(kind=kind)
+    elif task_name == "slcp":
+        return SLCP(kind=kind)
+    else:
+        raise ValueError(f"Unknown task: {task_name}")

gensbi_examples/utils.py

@@ -0,0 +1,15 @@
+# Utility function to get the checkpoint directory for a specific example
+import os
+
+# def get_example_checkpoint_dir(example_name):
+#     """
+#     Returns the absolute path to the checkpoint directory for a given example.
+#     Example:
+#         get_example_checkpoint_dir("my_first_model")
+#         # returns /path/to/GenSBI-examples/examples/getting_started/checkpoints
+#     """
+#     base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+#     if example_name == "my_first_model":
+#         return os.path.join(base_dir, "examples", "getting_started", "checkpoints")
+#     # Add more mappings as needed
+#     raise ValueError(f"Unknown example name: {example_name}")

gensbi_examples/utils.py.bk

@@ -0,0 +1,56 @@
+# Utility function to get the checkpoint directory for a specific example
+import os
+
+def get_example_checkpoint_dir(example_name):
+    """
+    Returns the absolute path to the checkpoint directory for a given example.
+    Example:
+        get_example_checkpoint_dir("my_first_model")
+        # returns /path/to/GenSBI-examples/examples/getting_started/checkpoints
+    """
+    base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    if example_name == "my_first_model":
+        return os.path.join(base_dir, "examples", "getting_started", "checkpoints")
+    # Add more mappings as needed
+    raise ValueError(f"Unknown example name: {example_name}")
+import os 
+import os
+from IPython import get_ipython
+
+def download_artifacts(task=None, dir=None):
+    """
+    Downloads the artifacts from the GenSBI repository.
+    """
+    root = "https://github.com/aurelio-amerio/GenSBI-examples/releases/download"
+    tag = "data-v0.1"
+    if task is not None:
+        fnames = [f"data_{task}.npz"]
+    else:
+        fnames =[
+            "data_two_moons.npz",
+            "data_bernoulli_glm.npz",
+            "data_gaussian_linear.npz",
+            "data_gaussian_linear_uniform.npz",
+            "data_gaussian_mixture.npz",
+            "data_slcp.npz"]
+        
+    fnames = [os.path.join(root, tag, fname) for fname in fnames]
+
+    if dir is None:
+        dir = os.path.join(os.getcwd(), "task_data")
+    else:
+        dir = os.path.join(dir, "task_data")
+    os.makedirs(dir, exist_ok=True)
+    for fname in fnames:
+        local_fname = os.path.join(dir, os.path.basename(fname))
+        if not os.path.exists(local_fname):
+            print(f"Downloading {fname} to {local_fname}")
+            os.system(f"wget -O {local_fname} {fname}")
+        else:
+            print(f"{local_fname} already exists, skipping download.")
+
+
+
+
+
+    

gensbi_examples-0.0.2.dist-info/METADATA

@@ -0,0 +1,72 @@
+Metadata-Version: 2.4
+Name: gensbi-examples
+Version: 0.0.2
+Summary: Examples for the GenSBI library
+Project-URL: Homepage, https://github.com/aurelio-amerio/GenSBI-examples
+Project-URL: Issues, https://github.com/aurelio-amerio/GenSBI-examples/issues
+Author-email: Aurelio Amerio <aure.amerio@gmail.com>
+License:    Copyright 2025 Amerio Aurelio
+        
+           Licensed under the Apache License, Version 2.0 (the "License");
+           you may not use this file except in compliance with the License.
+           You may obtain a copy of the License at
+        
+             http://www.apache.org/licenses/LICENSE-2.0
+        
+           Unless required by applicable law or agreed to in writing, software
+           distributed under the License is distributed on an "AS IS" BASIS,
+           WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+           See the License for the specific language governing permissions and
+           limitations under the License.
+License-File: LICENSE
+Requires-Python: >=3.11
+Requires-Dist: datasets
+Requires-Dist: flax>=0.12.0
+Requires-Dist: grain>=0.2.12
+Requires-Dist: huggingface-hub
+Requires-Dist: jax<=0.8.1,>=0.7.2
+Requires-Dist: matplotlib>=3.10
+Requires-Dist: numpy>=2.0
+Requires-Dist: scikit-learn>=1.7.0
+Description-Content-Type: text/markdown
+
+# GenSBI Examples
+
+This repository contains a collection of examples, tutorials, and recipes for **GenSBI**, a JAX-based library for Simulation-Based Inference using generative models.
+
+These examples demonstrate how to use GenSBI for various tasks, including:
+
+- Defining and running inference pipelines.
+- Using different embedding networks (MLP, ResNet, etc.).
+- Handling various data types (1D signals, 2D images).
+
+## Installation
+
+### Prerequisites
+
+You need to have **GenSBI** installed.
+
+**With CUDA 12 support (Recommended):**
+
+```bash
+pip install gensbi[cuda12]
+```
+
+**CPU-only:**
+
+```bash
+pip install gensbi
+```
+
+### Install Examples Package
+
+To run the examples and ensure all dependencies are met, install this package:
+
+```bash
+pip install gensbi-examples
+```
+
+## Structure
+
+- `examples/`: Contains standalone example scripts and notebooks.
+- `src/gensbi_examples`: Helper utilities for the examples.

gensbi_examples-0.0.2.dist-info/RECORD

@@ -0,0 +1,13 @@
+gensbi_examples/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+gensbi_examples/c2st.py,sha256=PU5fGq6QAItIyc_nscINk2TUrqvXYk-9PCDFVYbambg,2815
+gensbi_examples/c2st_v2.py.bk,sha256=xqkYZkUvmOnylBn0BmAY8n5kzc29OYu73U8EilBFn48,4455
+gensbi_examples/graph.py,sha256=OdJ8fiSP5_pgmzmHr2E9ZWYZzZm3ToJWLho2gRH0w7I,7570
+gensbi_examples/mask.py,sha256=OO7QH_r7SRjCiGLPn0vSS-mtFD2FE4JZFWhbed0MkcI,2997
+gensbi_examples/sbi_tasks.py.bk,sha256=AhJBXxBygXWSxmTzLzYGKrFI9hS9nuCHeLPZ5QgMWNk,14678
+gensbi_examples/tasks.py,sha256=-G20P3tQ4xwmjhJah1f0aHwdJXtNJUSgBlz2sGlcmaQ,10893
+gensbi_examples/utils.py,sha256=mF-wutjqqCX1EEmhHAtAlWvCSBcCiddoB9TBJjgz-xM,709
+gensbi_examples/utils.py.bk,sha256=_Q2hgbNlwo2uNFZ2tHIs_uye8IbLmsJYy6cYrErthT8,1863
+gensbi_examples-0.0.2.dist-info/METADATA,sha256=FsH7l_zy7WyFY5UCRm_aealQ6bvPW34yVpCIByu4Ctc,2233
+gensbi_examples-0.0.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+gensbi_examples-0.0.2.dist-info/licenses/LICENSE,sha256=D8Mi2-fbemv3oPZgMB-COT0aw2DXajsiebPYWtOMSpg,582
+gensbi_examples-0.0.2.dist-info/RECORD,,

gensbi_examples-0.0.2.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any

gensbi_examples-0.0.2.dist-info/licenses/LICENSE

@@ -0,0 +1,13 @@
+   Copyright 2025 Amerio Aurelio
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.