varselpfn 0.1.0__tar.gz

varselpfn-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Omar Melikechi
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

varselpfn-0.1.0/PKG-INFO

@@ -0,0 +1,67 @@
+Metadata-Version: 2.4
+Name: varselpfn
+Version: 0.1.0
+Summary: Variable selection with prior-data fitted networks: CD-PFN and PC-PFN
+Author-email: Omar Melikechi <omar.melikechi@duke.edu>
+License-Expression: MIT
+Project-URL: Repository, https://github.com/omelikechi/varselpfn
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+License-File: THIRD_PARTY_NOTICES
+Requires-Dist: numpy>=1.24
+Requires-Dist: safetensors>=0.4
+Requires-Dist: scipy>=1.10
+Requires-Dist: threadpoolctl>=3.0
+Requires-Dist: torch>=2.1
+Provides-Extra: pandas
+Requires-Dist: pandas>=1.5; extra == "pandas"
+Dynamic: license-file
+
+# Prior-data fitted networks for variable selection
+
+> Two pretrained prior-data fitted networks that return posterior inclusion probabilities for every covariate in one forward pass: **CD-PFN** estimates whether each covariate is conditionally dependent on the response, and **PC-PFN** estimates whether each covariate is a parent or child of the response in a directed acyclic graph.
+
+## Associated paper
+
+- **Prior-data fitted networks for variable selection**  
+	[arXiv](https://arxiv.org/abs/2611.05966)
+
+## Installation
+```
+pip install varselpfn
+```
+
+## Usage
+```python
+import numpy as np
+from varselpfn import CDPFN, PCPFN
+
+rng = np.random.default_rng(0)
+X = rng.normal(size=(200, 50))
+y = X[:, 0] + np.sin(2 * X[:, 1]) + rng.normal(size=200)
+
+# CD-PFN
+cd = CDPFN()
+pips = cd.predict(X, y)  # shape (50,)
+selected = cd.select(X, y, alpha=0.1)  # indices selected at target FDR 0.1
+
+# PC-PFN
+pc = PCPFN()
+pips = pc.predict(X, y)  # shape (50, 2): parent and child PIPs
+selected = pc.select(X, y, alpha=0.1)  # {'parents': [...], 'children': [...]}
+```
+
+- `X` is an `(n, p)` array or DataFrame of numeric covariates; `y` is continuous or binary (0/1). No preprocessing is needed, but missing values are not allowed.
+- The networks were trained on datasets with 16 to 1024 observations and 8 to 512 covariates.
+- Pass `device='cuda'` (or `'mps'`) to run on a GPU.
+
+## Fine-tuning
+
+To adapt a network to your data, copy its config ([`cd_pfn.json`](https://github.com/omelikechi/varselpfn/blob/main/varselpfn/configs/cd_pfn.json) or [`pc_pfn.json`](https://github.com/omelikechi/varselpfn/blob/main/varselpfn/configs/pc_pfn.json)), edit the `prior` block (options in [`prior.py`](https://github.com/omelikechi/varselpfn/blob/main/varselpfn/prior.py)), and continue training from the released weights:
+
+```
+python -m varselpfn.train --config my_prior.json --init cd_pfn --run-dir my_run --max-minutes 120
+```
+
+Then load it with `CDPFN(weights='my_run/best.pt')`. For PC-PFN, use `--init pc_pfn` and `PCPFN(weights=...)`.

varselpfn-0.1.0/README.md

@@ -0,0 +1,47 @@
+# Prior-data fitted networks for variable selection
+
+> Two pretrained prior-data fitted networks that return posterior inclusion probabilities for every covariate in one forward pass: **CD-PFN** estimates whether each covariate is conditionally dependent on the response, and **PC-PFN** estimates whether each covariate is a parent or child of the response in a directed acyclic graph.
+
+## Associated paper
+
+- **Prior-data fitted networks for variable selection**  
+	[arXiv](https://arxiv.org/abs/2611.05966)
+
+## Installation
+```
+pip install varselpfn
+```
+
+## Usage
+```python
+import numpy as np
+from varselpfn import CDPFN, PCPFN
+
+rng = np.random.default_rng(0)
+X = rng.normal(size=(200, 50))
+y = X[:, 0] + np.sin(2 * X[:, 1]) + rng.normal(size=200)
+
+# CD-PFN
+cd = CDPFN()
+pips = cd.predict(X, y)  # shape (50,)
+selected = cd.select(X, y, alpha=0.1)  # indices selected at target FDR 0.1
+
+# PC-PFN
+pc = PCPFN()
+pips = pc.predict(X, y)  # shape (50, 2): parent and child PIPs
+selected = pc.select(X, y, alpha=0.1)  # {'parents': [...], 'children': [...]}
+```
+
+- `X` is an `(n, p)` array or DataFrame of numeric covariates; `y` is continuous or binary (0/1). No preprocessing is needed, but missing values are not allowed.
+- The networks were trained on datasets with 16 to 1024 observations and 8 to 512 covariates.
+- Pass `device='cuda'` (or `'mps'`) to run on a GPU.
+
+## Fine-tuning
+
+To adapt a network to your data, copy its config ([`cd_pfn.json`](https://github.com/omelikechi/varselpfn/blob/main/varselpfn/configs/cd_pfn.json) or [`pc_pfn.json`](https://github.com/omelikechi/varselpfn/blob/main/varselpfn/configs/pc_pfn.json)), edit the `prior` block (options in [`prior.py`](https://github.com/omelikechi/varselpfn/blob/main/varselpfn/prior.py)), and continue training from the released weights:
+
+```
+python -m varselpfn.train --config my_prior.json --init cd_pfn --run-dir my_run --max-minutes 120
+```
+
+Then load it with `CDPFN(weights='my_run/best.pt')`. For PC-PFN, use `--init pc_pfn` and `PCPFN(weights=...)`.

varselpfn-0.1.0/THIRD_PARTY_NOTICES

@@ -0,0 +1,33 @@
+Third-party notices
+
+The Cauchy-DAG graph sampler in varselpfn/prior.py (sample_graph) is adapted from TabICL
+(https://github.com/soda-inria/tabicl, commit c91f00d), distributed under the following license.
+
+BSD 3-Clause License
+
+Copyright (c) 2025, Soda team @ Inria
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

varselpfn-0.1.0/pyproject.toml

@@ -0,0 +1,26 @@
+[build-system]
+requires = ['setuptools>=77']
+build-backend = 'setuptools.build_meta'
+
+[project]
+name = 'varselpfn'
+version = '0.1.0'
+description = 'Variable selection with prior-data fitted networks: CD-PFN and PC-PFN'
+readme = 'README.md'
+license = 'MIT'
+license-files = ['LICENSE', 'THIRD_PARTY_NOTICES']
+authors = [{name = 'Omar Melikechi', email = 'omar.melikechi@duke.edu'}]
+requires-python = '>=3.10'
+dependencies = ['numpy>=1.24', 'safetensors>=0.4', 'scipy>=1.10', 'threadpoolctl>=3.0', 'torch>=2.1']
+
+[project.optional-dependencies]
+pandas = ['pandas>=1.5']
+
+[project.urls]
+Repository = 'https://github.com/omelikechi/varselpfn'
+
+[tool.setuptools]
+packages = ['varselpfn']
+
+[tool.setuptools.package-data]
+varselpfn = ['configs/*.json', 'weights/*.safetensors']

varselpfn-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

varselpfn-0.1.0/varselpfn/__init__.py

@@ -0,0 +1,4 @@
+# varselpfn: variable selection with prior-data fitted networks
+from varselpfn.api import CDPFN, PCPFN, bayes_select
+
+__version__ = '0.1.0'

varselpfn-0.1.0/varselpfn/api.py

@@ -0,0 +1,131 @@
+# CD-PFN and PC-PFN
+'''
+CDPFN and PCPFN wrap the pretrained networks. predict(X, y) returns estimated posterior inclusion probabilities (PIPs)
+in one forward pass, with nothing fitted to the data; select(X, y, alpha) applies the Bayesian FDR rule to them.
+
+Weights are the bundled ones ('cd_pfn', 'pc_pfn') unless a path is given: a .safetensors file or a checkpoint written
+by varselpfn.train (best.pt), for example after fine-tuning.
+'''
+
+from functools import lru_cache
+import json
+from pathlib import Path
+
+import numpy as np
+from safetensors import safe_open
+from safetensors.torch import load_file
+import torch
+
+from varselpfn.model import Network
+
+package_dir = Path(__file__).resolve().parent
+bundled_weights = {'cd_pfn': package_dir / 'weights' / 'cd_pfn.safetensors', 'pc_pfn': package_dir / 'weights' / 'pc_pfn.safetensors'}
+
+#----------------------------------------------------------------
+# Weights
+#----------------------------------------------------------------
+def load_weights(weights):
+    '''Returns (model config, state dict) for a bundled name or a path to a .safetensors file or a training checkpoint.'''
+    path = bundled_weights.get(weights, Path(weights).expanduser())
+    if path.suffix == '.safetensors':
+        with safe_open(path, 'pt') as f:
+            model_config = json.loads(f.metadata()['model'])
+        return model_config, load_file(path)
+    checkpoint = torch.load(path, map_location='cpu', weights_only=False)
+    return checkpoint['config']['model'], checkpoint['model']
+
+@lru_cache(maxsize=None)
+def load_network(weights, device):
+    '''Loaded once per (weights, device) and reused.'''
+    model_config, state = load_weights(weights)
+    model = Network(**(model_config | {'grad_checkpoint': False}))
+    model.load_state_dict(state)
+    return model.to(device).eval()
+
+#----------------------------------------------------------------
+# Selection
+#----------------------------------------------------------------
+def bayes_select(pips, alpha=0.1):
+    '''
+    Bayesian FDR rule: the largest set of covariates with the highest PIPs whose mean (1 - PIP), the estimated FDR, is
+    at most alpha. Returns their indices, in decreasing order of PIP (ties by index).
+    '''
+    pips = np.asarray(pips, dtype=float)
+    p = len(pips)
+    order = np.lexsort((np.arange(p), -pips))
+    running = np.cumsum(1 - pips[order]) / np.arange(1, p + 1)
+    ok = np.flatnonzero(running <= alpha)
+    k = ok.max() + 1 if len(ok) else 0
+    return order[:k]
+
+#----------------------------------------------------------------
+# Models
+#----------------------------------------------------------------
+class PFN:
+    '''Shared by CDPFN and PCPFN: input checks, the forward pass, and covariate names for pandas input.'''
+    default_weights = None
+    outputs = None
+
+    def __init__(self, weights=None, device='cpu'):
+        self.weights = str(weights) if weights is not None else self.default_weights
+        self.device = torch.device(device)
+
+    def probabilities(self, X, y):
+        names = list(X.columns) if hasattr(X, 'columns') else None
+        X = np.asarray(X, dtype=np.float32)
+        y = np.asarray(y, dtype=np.float32).ravel()
+        if X.ndim != 2 or len(X) != len(y):
+            raise ValueError(f'X must be (n, p) and y (n,); got X {X.shape} and y {y.shape}')
+        if not (np.isfinite(X).all() and np.isfinite(y).all()):
+            raise ValueError('X and y must not contain missing or infinite values')
+        model = load_network(self.weights, self.device)
+        with torch.no_grad():
+            logits = model(torch.from_numpy(X)[None].to(self.device), torch.from_numpy(y)[None].to(self.device))[0]
+        pips = torch.sigmoid(logits[:, self.outputs].float()).cpu().numpy().astype(float)
+        return pips, names
+
+class CDPFN(PFN):
+    '''
+    CD-PFN: for each covariate, the estimated posterior probability that y depends on it given the other covariates.
+
+    X: (n, p) array or DataFrame of numeric covariates (continuous, binary, or ordinal). y: (n,) continuous or binary
+    response. Trained on 16 <= n <= 1024 and 8 <= p <= 512; columns are standardized inside the model.
+    '''
+    default_weights = 'cd_pfn'
+    outputs = 0
+
+    def predict(self, X, y):
+        '''PIPs, shape (p,); a Series indexed by the column names if X is a DataFrame.'''
+        pips, names = self.probabilities(X, y)
+        if names is None:
+            return pips
+        import pandas as pd
+        return pd.Series(pips, index=names, name='pip')
+
+    def select(self, X, y, alpha=0.1):
+        '''Covariates selected at target FDR alpha: indices, or column names if X is a DataFrame.'''
+        pips, names = self.probabilities(X, y)
+        selected = bayes_select(pips, alpha)
+        return selected if names is None else [names[j] for j in selected]
+
+class PCPFN(PFN):
+    '''
+    PC-PFN: for each covariate, the estimated posterior probabilities that it is a parent (direct cause) and a child
+    (direct effect) of y. Inputs as for CDPFN.
+    '''
+    default_weights = 'pc_pfn'
+    outputs = [1, 2]
+
+    def predict(self, X, y):
+        '''PIPs, shape (p, 2) with columns (parent, child); a DataFrame indexed by the column names if X is one.'''
+        pips, names = self.probabilities(X, y)
+        if names is None:
+            return pips
+        import pandas as pd
+        return pd.DataFrame(pips, index=names, columns=['parent', 'child'])
+
+    def select(self, X, y, alpha=0.1):
+        '''Parents and children selected at target FDR alpha, each by the Bayesian FDR rule: {'parents': ..., 'children': ...}.'''
+        pips, names = self.probabilities(X, y)
+        out = {'parents': bayes_select(pips[:, 0], alpha), 'children': bayes_select(pips[:, 1], alpha)}
+        return out if names is None else {role: [names[j] for j in selected] for role, selected in out.items()}

varselpfn-0.1.0/varselpfn/configs/cd_pfn.json

@@ -0,0 +1,174 @@
+{
+  "seed": 20261017,
+  "prior": {
+    "family": "general",
+    "p_range": [
+      8,
+      256
+    ],
+    "n_range": [
+      16,
+      1024
+    ],
+    "latent_zero_prob": 0.3,
+    "latent_fraction_range": [
+      0.0,
+      0.3
+    ],
+    "indegree_range": [
+      0.4,
+      2.5
+    ],
+    "cauchy_graph_prob": 0.5,
+    "mechanism_weights": {
+      "linear": 0.3,
+      "additive": 0.3,
+      "mlp": 0.12,
+      "tree": 0.12,
+      "gp": 0.08,
+      "product": 0.08
+    },
+    "additive_set": [
+      "identity",
+      "tanh",
+      "sin",
+      "cos",
+      "square",
+      "abs",
+      "bump",
+      "relu",
+      "softplus",
+      "sigmoid"
+    ],
+    "noise_families": [
+      "gaussian",
+      "laplace",
+      "uniform",
+      "student_t",
+      "mixture"
+    ],
+    "noise_fraction_range": [
+      0.05,
+      0.95
+    ],
+    "binary_feature_prob": 0.1,
+    "ordinal_feature_prob": 0.1,
+    "ordinal_levels": [
+      3,
+      7
+    ],
+    "binary_y_prob": 0.3,
+    "measurement_discretization_prob": 0.2,
+    "n_ref": 512,
+    "component_weights": {
+      "scm": 0.25,
+      "xfirst": 0.75
+    },
+    "xfirst": {
+      "structure_weights": {
+        "independent": 0.1,
+        "factor": 0.25,
+        "hierarchical": 0.35,
+        "ar": 0.1,
+        "scm": 0.2
+      },
+      "transform_prob": 0.7,
+      "transform_weights": {
+        "exp": 0.4,
+        "zero_inflated": 0.2,
+        "cube": 0.1,
+        "sqrt": 0.1,
+        "uniform": 0.1,
+        "bimodal": 0.1
+      },
+      "all_discrete_prob": 0.15,
+      "all_transform_prob": 0.25,
+      "continuous_prob": 0.4,
+      "null_prob": 0.1,
+      "max_parent_fraction": 0.5,
+      "children_prob": 0.3,
+      "max_children": 3,
+      "max_spouses": 2,
+      "spouses_near_child": true,
+      "descendant_prob": 0.0,
+      "max_descendants": 2,
+      "child_share": 0.0,
+      "parent_uniform_prob": 0.5,
+      "max_product_terms": 3
+    },
+    "p_uniform_prob": 0.5,
+    "n_uniform_prob": 0.5,
+    "additive_weights": {
+      "identity": 0.12,
+      "tanh": 0.12,
+      "sigmoid": 0.12,
+      "bump": 0.12,
+      "relu": 0.1,
+      "softplus": 0.1,
+      "square": 0.1,
+      "abs": 0.1,
+      "step": 0.08,
+      "sin": 0.02,
+      "cos": 0.02
+    },
+    "mlp_activations": [
+      "tanh",
+      "relu",
+      "sigmoid",
+      "bump",
+      "softplus",
+      "abs"
+    ],
+    "gp_length_range": [
+      1.0,
+      4.0
+    ],
+    "product_functions": [
+      "identity",
+      "tanh",
+      "sigmoid",
+      "bump",
+      "relu"
+    ],
+    "max_in_degree": 10,
+    "p_large_prob": 0.2,
+    "p_large_range": [
+      257,
+      512
+    ]
+  },
+  "model": {
+    "dim": 192,
+    "heads": 6,
+    "ff_dim": 768,
+    "blocks": 8,
+    "dropout": 0.0,
+    "pool": "mean_max",
+    "grad_checkpoint": false
+  },
+  "train": {
+    "cells_per_batch": 65536,
+    "max_batch_size": 64,
+    "lr": 0.0001,
+    "min_lr_ratio": 0.1,
+    "weight_decay": 0.0001,
+    "grad_clip": 1.0,
+    "warmup_steps": 2000,
+    "max_steps": 100000000,
+    "max_train_minutes": 1440,
+    "val_datasets": 512,
+    "val_every": 1000,
+    "log_every": 50,
+    "prevalence_datasets": 4000,
+    "workers": 15,
+    "schedule": "wsd",
+    "decay_fraction": 0.15,
+    "checkpoint_cells": 150000,
+    "head_weights": [
+      1.0,
+      0.0,
+      0.0
+    ],
+    "val_match_training": true
+  }
+}