gensbi-examples 0.0.2__tar.gz

gensbi_examples-0.0.2/.gitignore

@@ -0,0 +1,217 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+#poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+#pdm.lock
+#pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+#pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  you could uncomment the following to ignore the entire vscode folder
+.vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Cursor
+#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
+#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
+#  refer to https://docs.cursor.com/context/ignore-files
+.cursorignore
+.cursorindexingignore
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# task data
+task_data/*
+*.npz
+
+#sub
+condor_logs/
+
+#some test traces
+tests/tmp/*

gensbi_examples-0.0.2/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.

gensbi_examples-0.0.2/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,40 @@
+# 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/gensbi_examples/c2st.py

@@ -0,0 +1,111 @@
+from typing import Optional
+
+import jax
+from jax import numpy as jnp
+from jax import Array
+import numpy as np
+from sklearn.model_selection import KFold, cross_val_score
+from sklearn.neural_network import MLPClassifier
+
+
+def c2st(
+    X: Array,
+    Y: Array,
+    seed: int = 1,
+    n_folds: int = 5,
+    scoring: str = "accuracy",
+    z_score: bool = True,
+    noise_scale: Optional[float] = None,
+) -> Array:
+    """Classifier-based 2-sample test returning accuracy
+
+    Trains classifiers with N-fold cross-validation [1]. Scikit learn MLPClassifier are
+    used, with 2 hidden layers of 10x dim each, where dim is the dimensionality of the
+    samples X and Y.
+
+    Args:
+        X: Sample 1
+        Y: Sample 2
+        seed: Seed for sklearn
+        n_folds: Number of folds
+        z_score: Z-scoring using X
+        noise_scale: If passed, will add Gaussian noise with std noise_scale to samples
+
+    References:
+        [1]: https://scikit-learn.org/stable/modules/cross_validation.html
+    """
+    if z_score:
+        X_mean = jnp.mean(X, axis=0)
+        X_std = jnp.std(X, axis=0)
+        X = (X - X_mean) / X_std
+        Y = (Y - X_mean) / X_std
+
+
+    if noise_scale is not None:
+        key = jax.random.PRNGKey(seed)
+        X += noise_scale * jax.random.normal(key, X.shape) * noise_scale
+        Y += noise_scale * jax.random.normal(key, Y.shape) * noise_scale
+
+
+    # Convert to numpy if not already
+
+    X = np.asarray(X)
+    Y = np.asarray(Y)
+
+    ndim = X.shape[1]
+
+    clf = MLPClassifier(
+        activation="relu",
+        hidden_layer_sizes=(10 * ndim, 10 * ndim),
+        max_iter=10000,
+        solver="adam",
+        random_state=seed,
+    )
+
+    data = np.concatenate((X, Y))
+    target = np.concatenate(
+        (
+            np.zeros((X.shape[0],)),
+            np.ones((Y.shape[0],)),
+        )
+    )
+
+    shuffle = KFold(n_splits=n_folds, shuffle=True, random_state=seed)
+    scores = cross_val_score(clf, data, target, cv=shuffle, scoring=scoring)
+
+    scores = np.asarray(np.mean(scores)).astype(np.float32)
+    return scores
+
+
+def c2st_auc(
+    X: Array,
+    Y: Array,
+    seed: int = 1,
+    n_folds: int = 5,
+    z_score: bool = True,
+    noise_scale: Optional[float] = None,
+    ) -> Array:
+    """Classifier-based 2-sample test returning AUC (area under curve)
+
+    Same as c2st, except that it returns ROC AUC rather than accuracy
+
+    Args:
+        X: Sample 1
+        Y: Sample 2
+        seed: Seed for sklearn
+        n_folds: Number of folds
+        z_score: Z-scoring using X
+        noise_scale: If passed, will add Gaussian noise with std noise_scale to samples
+
+    Returns:
+        Metric
+    """
+    return c2st(
+        X,
+        Y,
+        seed=seed,
+        n_folds=n_folds,
+        scoring="roc_auc",
+        z_score=z_score,
+        noise_scale=noise_scale,
+    )

gensbi_examples-0.0.2/gensbi_examples/c2st_v2.py.bk

@@ -0,0 +1,147 @@
+from typing import Optional
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+import numpy as np
+from sklearn.model_selection import KFold
+from flax import nnx
+import optax
+
+# Define MLP using flax.nnx
+class MLP(nnx.Module):
+    def __init__(self, in_dim, hidden_dim, *, rngs):
+        self.seq = nnx.Sequential([
+            nnx.Linear(in_dim, hidden_dim, rngs= rngs),
+            nnx.Relu(),
+            nnx.Linear(hidden_dim, hidden_dim, rngs=rngs),
+            nnx.Relu(),
+            nnx.Linear(hidden_dim, 2, rngs=rngs),
+        ])
+    def __call__(self, x):
+        return self.seq(x)
+
+
+def loss_fn(state, x, y):
+    logits = state.value(x)
+    labels = jax.nn.one_hot(y, 2)
+    loss = optax.softmax_cross_entropy(logits, labels).mean()
+    return loss
+
+def accuracy_fn(state, x, y):
+    logits = state.value(x)
+    preds = jnp.argmax(logits, axis=-1)
+    return (preds == y).mean()
+
+def c2st(
+    X: Array,
+    Y: Array,
+    seed: int = 1,
+    n_folds: int = 5,
+    z_score: bool = True,
+    noise_scale: Optional[float] = None,
+) -> Array:
+    """Classifier-based 2-sample test returning accuracy (using nnx for GPU training)
+
+    Trains classifiers with N-fold cross-validation [1]. nnx MLP is used, with 2 hidden layers of 10x dim each.
+
+    Args:
+        X: Sample 1
+        Y: Sample 2
+        seed: Seed for random number generation
+        n_folds: Number of folds
+        z_score: Z-scoring using X
+        noise_scale: If passed, will add Gaussian noise with std noise_scale to samples
+
+    References:
+        [1]: https://scikit-learn.org/stable/modules/cross_validation.html
+    """
+
+    rngs = nnx.Rngs(seed)
+    if z_score:
+        X_mean = jnp.mean(X, axis=0)
+        X_std = jnp.std(X, axis=0)
+        X = (X - X_mean) / X_std
+        Y = (Y - X_mean) / X_std
+
+    if noise_scale is not None:
+        key = jax.random.PRNGKey(seed)
+        X = X + noise_scale * jax.random.normal(key, X.shape)
+        Y = Y + noise_scale * jax.random.normal(key, Y.shape)
+
+    X = jnp.asarray(X)
+    Y = jnp.asarray(Y)
+    ndim = X.shape[1]
+
+    # Prepare data and targets
+    data = jnp.concatenate([X, Y], axis=0)
+    target = jnp.concatenate([
+        jnp.zeros((X.shape[0],), dtype=jnp.int32),
+        jnp.ones((Y.shape[0],), dtype=jnp.int32)
+    ], axis=0)
+
+    kf = KFold(n_splits=n_folds, shuffle=True, random_state=seed)
+    scores = []
+
+    for fold, (train_idx, test_idx) in enumerate(kf.split(data)):
+        x_train, y_train = data[train_idx], target[train_idx]
+        x_test, y_test = data[test_idx], target[test_idx]
+
+    for fold, (train_idx, test_idx) in enumerate(kf.split(data)):
+        # Model and optimizer
+        key = jax.random.PRNGKey(seed + fold)
+        model = MLP(ndim, 10 * ndim, rngs=rngs)
+        optimizer = nnx.Optimizer(model, optax.adam(1e-3))
+
+        @jax.jit
+        def train_step(optimizer, x, y):
+            def _loss_fn(model):
+                return loss_fn(model, x, y)
+            loss, grads = nnx.value_and_grad(_loss_fn)(optimizer.target)
+            optimizer.update(grads, value=loss)
+            return optimizer, loss
+
+        # Training loop
+        n_epochs = 100
+        batch_size = min(128, x_train.shape[0])
+        n_batches = int(jnp.ceil(x_train.shape[0] / batch_size))
+        for epoch in range(n_epochs):
+            perm = jax.random.permutation(key, x_train.shape[0])
+            x_train_shuffled = x_train[perm]
+            y_train_shuffled = y_train[perm]
+            for i in range(n_batches):
+                start = i * batch_size
+                end = min((i + 1) * batch_size, x_train.shape[0])
+                xb = x_train_shuffled[start:end]
+                yb = y_train_shuffled[start:end]
+                optimizer, _ = train_step(optimizer, xb, yb)
+        model = optimizer.target
+
+        score = float(accuracy_fn(model, x_test, y_test))
+
+        scores.append(score)
+
+    return np.asarray(np.mean(scores), dtype=np.float32)
+
+
+# def c2st_auc(
+#     X: Array,
+#     Y: Array,
+#     seed: int = 1,
+#     n_folds: int = 5,
+#     z_score: bool = True,
+#     noise_scale: Optional[float] = None,
+# ) -> Array:
+#     """Classifier-based 2-sample test returning AUC (area under curve)
+
+#     Same as c2st, except that it returns ROC AUC rather than accuracy
+#     """
+#     return c2st(
+#         X,
+#         Y,
+#         seed=seed,
+#         n_folds=n_folds,
+#         scoring="roc_auc",
+#         z_score=z_score,
+#         noise_scale=noise_scale,
+#     )