cpfn 1.0.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- cpfn-1.0.0/LICENCE +19 -0
- cpfn-1.0.0/PKG-INFO +127 -0
- cpfn-1.0.0/README.md +93 -0
- cpfn-1.0.0/pyproject.toml +28 -0
- cpfn-1.0.0/setup.cfg +4 -0
- cpfn-1.0.0/src/cpfn/__init__.py +4 -0
- cpfn-1.0.0/src/cpfn/estimator.py +288 -0
- cpfn-1.0.0/src/cpfn.egg-info/PKG-INFO +127 -0
- cpfn-1.0.0/src/cpfn.egg-info/SOURCES.txt +11 -0
- cpfn-1.0.0/src/cpfn.egg-info/dependency_links.txt +1 -0
- cpfn-1.0.0/src/cpfn.egg-info/requires.txt +3 -0
- cpfn-1.0.0/src/cpfn.egg-info/top_level.txt +1 -0
- cpfn-1.0.0/tests/test_estimator.py +26 -0
cpfn-1.0.0/LICENCE
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
Copyright (c) 2018 The Python Packaging Authority
|
|
2
|
+
|
|
3
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
4
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
5
|
+
in the Software without restriction, including without limitation the rights
|
|
6
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
7
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
8
|
+
furnished to do so, subject to the following conditions:
|
|
9
|
+
|
|
10
|
+
The above copyright notice and this permission notice shall be included in all
|
|
11
|
+
copies or substantial portions of the Software.
|
|
12
|
+
|
|
13
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
14
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
15
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
16
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
17
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
18
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
19
|
+
SOFTWARE.
|
cpfn-1.0.0/PKG-INFO
ADDED
|
@@ -0,0 +1,127 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: cpfn
|
|
3
|
+
Version: 1.0.0
|
|
4
|
+
Summary: Conditional Push-Forward Neural Network estimator
|
|
5
|
+
Author-email: tedescolor <tedescolor@gmail.com>
|
|
6
|
+
License: Copyright (c) 2018 The Python Packaging Authority
|
|
7
|
+
|
|
8
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
9
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
10
|
+
in the Software without restriction, including without limitation the rights
|
|
11
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
12
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
13
|
+
furnished to do so, subject to the following conditions:
|
|
14
|
+
|
|
15
|
+
The above copyright notice and this permission notice shall be included in all
|
|
16
|
+
copies or substantial portions of the Software.
|
|
17
|
+
|
|
18
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
19
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
20
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
21
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
22
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
23
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
24
|
+
SOFTWARE.
|
|
25
|
+
Project-URL: Repository, https://github.com/tedescolor/cpfn
|
|
26
|
+
Project-URL: Paper, https://arxiv.org/pdf/2511.14455
|
|
27
|
+
Requires-Python: >=3.8
|
|
28
|
+
Description-Content-Type: text/markdown
|
|
29
|
+
License-File: LICENCE
|
|
30
|
+
Requires-Dist: torch>=2.0
|
|
31
|
+
Requires-Dist: numpy>=1.20
|
|
32
|
+
Requires-Dist: tqdm>=4.60
|
|
33
|
+
Dynamic: license-file
|
|
34
|
+
|
|
35
|
+
# CPFN — Conditional Push-Forward Neural Network
|
|
36
|
+
|
|
37
|
+
Compact, importable implementation of a Conditional Push-Forward Neural Network (CPFN) estimator.
|
|
38
|
+
|
|
39
|
+
**Paper:** https://arxiv.org/pdf/2511.14455
|
|
40
|
+
|
|
41
|
+
## Goals
|
|
42
|
+
- Provide a lightweight `CPFN` class for estimating conditional generators.
|
|
43
|
+
- Expose a simple API for training and sampling.
|
|
44
|
+
|
|
45
|
+
## Install
|
|
46
|
+
|
|
47
|
+
### From PyPI
|
|
48
|
+
```bash
|
|
49
|
+
pip install cpfn
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
## Quick Usage
|
|
53
|
+
```python
|
|
54
|
+
import torch
|
|
55
|
+
import matplotlib.pyplot as plt
|
|
56
|
+
import numpy as np
|
|
57
|
+
import random
|
|
58
|
+
from cpfn import CPFN
|
|
59
|
+
|
|
60
|
+
# 0. Hardware Selection (CUDA for NVIDIA, MPS for Apple Silicon, or CPU) and Reproducibility
|
|
61
|
+
if torch.cuda.is_available():
|
|
62
|
+
device = torch.device("cuda")
|
|
63
|
+
elif torch.backends.mps.is_available():
|
|
64
|
+
device = torch.device("mps")
|
|
65
|
+
else:
|
|
66
|
+
device = torch.device("cpu")
|
|
67
|
+
|
|
68
|
+
SEED = 43
|
|
69
|
+
|
|
70
|
+
random.seed(SEED)
|
|
71
|
+
np.random.seed(SEED)
|
|
72
|
+
torch.manual_seed(SEED)
|
|
73
|
+
|
|
74
|
+
# 1. Define Synthetic Ground Truth (Branching Function)
|
|
75
|
+
def true_sample(x):
|
|
76
|
+
z = np.random.randn() # Gaussian noise
|
|
77
|
+
p = np.random.rand() # Uniform switch
|
|
78
|
+
|
|
79
|
+
# Conditional logic: creates two paths for x > 0.5
|
|
80
|
+
if x < 0.5 or p < 0.5:
|
|
81
|
+
return 10 * x * (x - 0.5) * (1.5 - x) + z * 0.3 * (1.3 - x)
|
|
82
|
+
else:
|
|
83
|
+
return 10 * x * (x - 0.5) * (0.8 - x) + z * 0.3 * (1.3 - x)
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
# 2. Generate Training Data
|
|
87
|
+
ntrain = 500
|
|
88
|
+
xs = np.random.rand(ntrain)
|
|
89
|
+
ys = np.array([true_sample(x) for x in xs])
|
|
90
|
+
|
|
91
|
+
# 3. Model Setup & Training
|
|
92
|
+
model = CPFN(d=1, q=1, r=20, width=50, hidden_layers=3, delta=1e-15)
|
|
93
|
+
model.to(device)
|
|
94
|
+
model.fit(xs, ys,
|
|
95
|
+
epochs=3000,
|
|
96
|
+
lr=1e-3,
|
|
97
|
+
m=30,
|
|
98
|
+
h0=5e-2)
|
|
99
|
+
|
|
100
|
+
model.freeze()
|
|
101
|
+
|
|
102
|
+
# 4. Inference: Generate 1 sample for every x in training set
|
|
103
|
+
# samples shape: (ntrain, 1, 1)
|
|
104
|
+
ys_gen = model.sample_conditional(xs, num_samples=1).flatten()
|
|
105
|
+
|
|
106
|
+
# 5. Visualization
|
|
107
|
+
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5), sharey=True)
|
|
108
|
+
|
|
109
|
+
ax1.scatter(xs, ys, alpha=0.6, s=15, label="Ground Truth")
|
|
110
|
+
ax1.set_title("Original Training Data")
|
|
111
|
+
ax1.set_xlabel("x")
|
|
112
|
+
ax1.set_ylabel("y")
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
## Tests
|
|
117
|
+
Run the included pytest smoke test:
|
|
118
|
+
```bash
|
|
119
|
+
pytest -q
|
|
120
|
+
```
|
|
121
|
+
|
|
122
|
+
## Development
|
|
123
|
+
- Source: `src/cpfn/`
|
|
124
|
+
- Tests: `tests/`
|
|
125
|
+
|
|
126
|
+
## License
|
|
127
|
+
See `LICENCE` in the repository root.
|
cpfn-1.0.0/README.md
ADDED
|
@@ -0,0 +1,93 @@
|
|
|
1
|
+
# CPFN — Conditional Push-Forward Neural Network
|
|
2
|
+
|
|
3
|
+
Compact, importable implementation of a Conditional Push-Forward Neural Network (CPFN) estimator.
|
|
4
|
+
|
|
5
|
+
**Paper:** https://arxiv.org/pdf/2511.14455
|
|
6
|
+
|
|
7
|
+
## Goals
|
|
8
|
+
- Provide a lightweight `CPFN` class for estimating conditional generators.
|
|
9
|
+
- Expose a simple API for training and sampling.
|
|
10
|
+
|
|
11
|
+
## Install
|
|
12
|
+
|
|
13
|
+
### From PyPI
|
|
14
|
+
```bash
|
|
15
|
+
pip install cpfn
|
|
16
|
+
```
|
|
17
|
+
|
|
18
|
+
## Quick Usage
|
|
19
|
+
```python
|
|
20
|
+
import torch
|
|
21
|
+
import matplotlib.pyplot as plt
|
|
22
|
+
import numpy as np
|
|
23
|
+
import random
|
|
24
|
+
from cpfn import CPFN
|
|
25
|
+
|
|
26
|
+
# 0. Hardware Selection (CUDA for NVIDIA, MPS for Apple Silicon, or CPU) and Reproducibility
|
|
27
|
+
if torch.cuda.is_available():
|
|
28
|
+
device = torch.device("cuda")
|
|
29
|
+
elif torch.backends.mps.is_available():
|
|
30
|
+
device = torch.device("mps")
|
|
31
|
+
else:
|
|
32
|
+
device = torch.device("cpu")
|
|
33
|
+
|
|
34
|
+
SEED = 43
|
|
35
|
+
|
|
36
|
+
random.seed(SEED)
|
|
37
|
+
np.random.seed(SEED)
|
|
38
|
+
torch.manual_seed(SEED)
|
|
39
|
+
|
|
40
|
+
# 1. Define Synthetic Ground Truth (Branching Function)
|
|
41
|
+
def true_sample(x):
|
|
42
|
+
z = np.random.randn() # Gaussian noise
|
|
43
|
+
p = np.random.rand() # Uniform switch
|
|
44
|
+
|
|
45
|
+
# Conditional logic: creates two paths for x > 0.5
|
|
46
|
+
if x < 0.5 or p < 0.5:
|
|
47
|
+
return 10 * x * (x - 0.5) * (1.5 - x) + z * 0.3 * (1.3 - x)
|
|
48
|
+
else:
|
|
49
|
+
return 10 * x * (x - 0.5) * (0.8 - x) + z * 0.3 * (1.3 - x)
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
# 2. Generate Training Data
|
|
53
|
+
ntrain = 500
|
|
54
|
+
xs = np.random.rand(ntrain)
|
|
55
|
+
ys = np.array([true_sample(x) for x in xs])
|
|
56
|
+
|
|
57
|
+
# 3. Model Setup & Training
|
|
58
|
+
model = CPFN(d=1, q=1, r=20, width=50, hidden_layers=3, delta=1e-15)
|
|
59
|
+
model.to(device)
|
|
60
|
+
model.fit(xs, ys,
|
|
61
|
+
epochs=3000,
|
|
62
|
+
lr=1e-3,
|
|
63
|
+
m=30,
|
|
64
|
+
h0=5e-2)
|
|
65
|
+
|
|
66
|
+
model.freeze()
|
|
67
|
+
|
|
68
|
+
# 4. Inference: Generate 1 sample for every x in training set
|
|
69
|
+
# samples shape: (ntrain, 1, 1)
|
|
70
|
+
ys_gen = model.sample_conditional(xs, num_samples=1).flatten()
|
|
71
|
+
|
|
72
|
+
# 5. Visualization
|
|
73
|
+
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5), sharey=True)
|
|
74
|
+
|
|
75
|
+
ax1.scatter(xs, ys, alpha=0.6, s=15, label="Ground Truth")
|
|
76
|
+
ax1.set_title("Original Training Data")
|
|
77
|
+
ax1.set_xlabel("x")
|
|
78
|
+
ax1.set_ylabel("y")
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
## Tests
|
|
83
|
+
Run the included pytest smoke test:
|
|
84
|
+
```bash
|
|
85
|
+
pytest -q
|
|
86
|
+
```
|
|
87
|
+
|
|
88
|
+
## Development
|
|
89
|
+
- Source: `src/cpfn/`
|
|
90
|
+
- Tests: `tests/`
|
|
91
|
+
|
|
92
|
+
## License
|
|
93
|
+
See `LICENCE` in the repository root.
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "cpfn"
|
|
3
|
+
version = "1.0.0"
|
|
4
|
+
description = "Conditional Push-Forward Neural Network estimator"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.8"
|
|
7
|
+
license = { file = "LICENCE" }
|
|
8
|
+
authors = [ { name = "tedescolor", email = "tedescolor@gmail.com" } ]
|
|
9
|
+
dependencies = [
|
|
10
|
+
"torch>=2.0",
|
|
11
|
+
"numpy>=1.20",
|
|
12
|
+
"tqdm>=4.60",
|
|
13
|
+
]
|
|
14
|
+
|
|
15
|
+
[project.urls]
|
|
16
|
+
"Repository" = "https://github.com/tedescolor/cpfn"
|
|
17
|
+
"Paper" = "https://arxiv.org/pdf/2511.14455"
|
|
18
|
+
|
|
19
|
+
[tool.setuptools]
|
|
20
|
+
packages = ["cpfn"]
|
|
21
|
+
|
|
22
|
+
[tool.setuptools.package-dir]
|
|
23
|
+
"" = "src"
|
|
24
|
+
|
|
25
|
+
[build-system]
|
|
26
|
+
requires = ["setuptools >= 77.0.3"]
|
|
27
|
+
build-backend = "setuptools.build_meta"
|
|
28
|
+
|
cpfn-1.0.0/setup.cfg
ADDED
|
@@ -0,0 +1,288 @@
|
|
|
1
|
+
import math
|
|
2
|
+
from typing import Optional
|
|
3
|
+
|
|
4
|
+
import torch
|
|
5
|
+
import torch.nn as nn
|
|
6
|
+
import torch.nn.functional as F
|
|
7
|
+
from tqdm import tqdm
|
|
8
|
+
import copy
|
|
9
|
+
|
|
10
|
+
gelu = nn.GELU()
|
|
11
|
+
|
|
12
|
+
class MLP(nn.Module):
|
|
13
|
+
def __init__(self, in_dim: int, out_dim: int, hidden_width: int = 50, hidden_layers: int = 3, activation: nn.Module = gelu, final_activation: bool = False):
|
|
14
|
+
super().__init__()
|
|
15
|
+
layers = []
|
|
16
|
+
d = in_dim
|
|
17
|
+
# Dynamically build hidden layers
|
|
18
|
+
for _ in range(hidden_layers):
|
|
19
|
+
layers.append(nn.Linear(d, hidden_width))
|
|
20
|
+
layers.append(activation)
|
|
21
|
+
d = hidden_width
|
|
22
|
+
# Output layer
|
|
23
|
+
layers.append(nn.Linear(d, out_dim))
|
|
24
|
+
if(final_activation):
|
|
25
|
+
layers.append(activation)
|
|
26
|
+
self.net = nn.Sequential(*layers)
|
|
27
|
+
self.final_activation = final_activation
|
|
28
|
+
|
|
29
|
+
def forward(self, x: torch.Tensor):
|
|
30
|
+
return self.net(x)
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
class CPFN(nn.Module):
|
|
34
|
+
"""Compact CPFN estimator suitable for import and use.
|
|
35
|
+
|
|
36
|
+
Usage: from cpfn.estimator import CPFN
|
|
37
|
+
model = CPFN(d=1, q=1)
|
|
38
|
+
model.to(device)
|
|
39
|
+
model.fit(xs, ys, ...)
|
|
40
|
+
"""
|
|
41
|
+
def __init__(
|
|
42
|
+
self,
|
|
43
|
+
d: int, # Input dimension (x)
|
|
44
|
+
q: int, # Output dimension (y)
|
|
45
|
+
r: int = 20, # Rank (latent factor dimension)
|
|
46
|
+
width: int = 50,
|
|
47
|
+
hidden_layers: int = 3,
|
|
48
|
+
latent_dist: str = "normal", # Distribution for latent variable u
|
|
49
|
+
learn_eps: bool = True, # Whether to learn the bandwidth parameter
|
|
50
|
+
eps_init: float = 5e-2,
|
|
51
|
+
delta: float = 1e-15, # Numerical stability constant
|
|
52
|
+
psi_final_activation: bool = True,
|
|
53
|
+
|
|
54
|
+
):
|
|
55
|
+
super().__init__()
|
|
56
|
+
self.d, self.q, self.r = d, q, r
|
|
57
|
+
self.latent_dist = latent_dist
|
|
58
|
+
self.delta = float(delta)
|
|
59
|
+
self._istraining = False
|
|
60
|
+
|
|
61
|
+
# Neural networks for the conditional (varphi) and latent (psi) components
|
|
62
|
+
# Outputs are flattened matrices of shape (r * q)
|
|
63
|
+
self.varphi = MLP(d, r * q, hidden_width=width, hidden_layers=hidden_layers, final_activation=False)
|
|
64
|
+
self.psi = MLP(q, r * q, hidden_width=width, hidden_layers=hidden_layers, final_activation=psi_final_activation)
|
|
65
|
+
|
|
66
|
+
# We use exp() to ensuring positivity
|
|
67
|
+
eps0 = torch.tensor([eps_init] * q, dtype=torch.float32)
|
|
68
|
+
if learn_eps:
|
|
69
|
+
self._eps_param = nn.Parameter(torch.log(eps0))
|
|
70
|
+
else:
|
|
71
|
+
self.register_buffer("_eps_param", torch.log(eps0))
|
|
72
|
+
self._eps_param.requires_grad_(False)
|
|
73
|
+
|
|
74
|
+
# Standardization statistics (initialized as None, computed during fit)
|
|
75
|
+
self.register_buffer("x_mean", None)
|
|
76
|
+
self.register_buffer("x_std", None)
|
|
77
|
+
self.register_buffer("y_mean", None)
|
|
78
|
+
self.register_buffer("y_std", None)
|
|
79
|
+
|
|
80
|
+
def eps(self) -> torch.Tensor:
|
|
81
|
+
# CHANGE 2: Use exp() to ensure positivity (like dlroms with direct parameter)
|
|
82
|
+
return torch.exp(self._eps_param) + 1e-12
|
|
83
|
+
|
|
84
|
+
def _sample_u(self, n: int, m: int, device: torch.device) -> torch.Tensor:
|
|
85
|
+
# Sample latent variable u from defined prior
|
|
86
|
+
if self.latent_dist == "uniform":
|
|
87
|
+
return torch.rand(n, m, self.q, device=device)
|
|
88
|
+
return torch.randn(n, m, self.q, device=device)
|
|
89
|
+
|
|
90
|
+
def _standardize_x(self, x: torch.Tensor) -> torch.Tensor:
|
|
91
|
+
"""Standardize input x using stored statistics."""
|
|
92
|
+
if self.x_mean is None or self.x_std is None:
|
|
93
|
+
return x
|
|
94
|
+
return (x - self.x_mean) / self.x_std
|
|
95
|
+
|
|
96
|
+
def _standardize_y(self, y: torch.Tensor) -> torch.Tensor:
|
|
97
|
+
"""Standardize output y using stored statistics."""
|
|
98
|
+
if self.y_mean is None or self.y_std is None:
|
|
99
|
+
return y
|
|
100
|
+
return (y - self.y_mean) / self.y_std
|
|
101
|
+
|
|
102
|
+
def _destandardize_y(self, y: torch.Tensor) -> torch.Tensor:
|
|
103
|
+
"""Convert standardized y back to original scale."""
|
|
104
|
+
if self.y_mean is None or self.y_std is None:
|
|
105
|
+
return y
|
|
106
|
+
return y * self.y_std + self.y_mean
|
|
107
|
+
|
|
108
|
+
def forward(self, x: torch.Tensor, u: torch.Tensor) -> torch.Tensor:
|
|
109
|
+
# Case 1: u is 2D (Batch, Q) - Single sample per input
|
|
110
|
+
if u.dim() == 2:
|
|
111
|
+
# Reshape output to (Batch, Rank, OutputDim)
|
|
112
|
+
vx = self.varphi(x).reshape(-1, self.r, self.q)
|
|
113
|
+
vu = self.psi(u).reshape(-1, self.r, self.q)
|
|
114
|
+
# Dot product interaction along the rank dimension
|
|
115
|
+
return (vx * vu).sum(dim=1)
|
|
116
|
+
|
|
117
|
+
# Case 2: u is 3D (Batch, M_samples, Q) - Multiple samples per input
|
|
118
|
+
if u.dim() == 3:
|
|
119
|
+
n, m, q = u.shape
|
|
120
|
+
vx = self.varphi(x).reshape(n, self.r, self.q)
|
|
121
|
+
# Process flattened u, then reshape back to (Batch, M_samples, Rank, Q)
|
|
122
|
+
vu = self.psi(u.reshape(n * m, q)).reshape(n, m, self.r, self.q)
|
|
123
|
+
# Broadcasting vx to match m samples: (Batch, 1, Rank, Q) * (Batch, M, Rank, Q)
|
|
124
|
+
y = (vx[:, None, :, :] * vu).sum(dim=2)
|
|
125
|
+
return y
|
|
126
|
+
|
|
127
|
+
raise ValueError("u must have shape (n,q) or (n,m,q).")
|
|
128
|
+
|
|
129
|
+
def logdensity(self, xs: torch.Tensor, ys : torch.Tensor, m : int = 30, tilted : bool = False):
|
|
130
|
+
# Check if inputs are tensors (or if training), otherwise delegate to numpy handler
|
|
131
|
+
if(self._istraining or (isinstance(xs, torch.Tensor) and isinstance(ys, torch.Tensor))):
|
|
132
|
+
xs_, ys_ = xs.reshape(-1, self.d), ys.reshape(-1, self.q)
|
|
133
|
+
|
|
134
|
+
# Standardize inputs and outputs
|
|
135
|
+
xs_ = self._standardize_x(xs_)
|
|
136
|
+
ys_ = self._standardize_y(ys_)
|
|
137
|
+
|
|
138
|
+
delta = self.delta if tilted else 1e-15
|
|
139
|
+
# Sample m latent variables per input
|
|
140
|
+
u = self._sample_u(xs_.shape[0], m, device=xs.device)
|
|
141
|
+
# Predict yhat based on x and random u
|
|
142
|
+
yhat = self.forward(xs_, u)
|
|
143
|
+
|
|
144
|
+
# Calculate residuals between actual y and generated hypotheses
|
|
145
|
+
residuals = (ys_[:, None, :] - yhat)
|
|
146
|
+
eps = self.eps()
|
|
147
|
+
|
|
148
|
+
# Compute Gaussian log-likelihood components
|
|
149
|
+
zs = residuals / eps
|
|
150
|
+
rs = zs.pow(2).sum(dim=-1) # Squared Mahalanobis-like distance
|
|
151
|
+
|
|
152
|
+
# Log-pdf of the Gaussian kernel
|
|
153
|
+
exponents = -0.5 * rs - 0.5 * self.q * math.log(2.0 * math.pi) - torch.log(eps).sum() - math.log(m) - math.log(delta)
|
|
154
|
+
|
|
155
|
+
# Prepare for LogSumExp
|
|
156
|
+
shape = list(exponents.shape)
|
|
157
|
+
shape[1] = 1
|
|
158
|
+
# Compute log mean exp (using stable logsumexp trick) over m samples
|
|
159
|
+
# Includes a stability term (zeros) to prevent -inf
|
|
160
|
+
logd = torch.logsumexp(torch.cat([exponents, torch.zeros(*shape, device = xs.device)], dim=1), dim=1) + math.log(delta)
|
|
161
|
+
|
|
162
|
+
# Return scalar if input was single item, else tensor
|
|
163
|
+
return logd if(len(xs.shape)==2 or xs_.shape[0]>1 or len(ys.shape)==2 or ys_.shape[0]>1) else logd[0]
|
|
164
|
+
else:
|
|
165
|
+
# Handle Numpy inputs by converting to Tensor and recurring
|
|
166
|
+
device = self.eps().device
|
|
167
|
+
return self.logdensity(torch.tensor(xs, device=device, dtype=torch.float32), torch.tensor(ys, device=device, dtype=torch.float32), m = m, tilted = tilted).cpu().numpy()
|
|
168
|
+
|
|
169
|
+
def sample_conditional(self, x: torch.Tensor, num_samples: int = 1, seed: Optional[int] = None) -> torch.Tensor:
|
|
170
|
+
if(self._istraining or isinstance(x, torch.Tensor)):
|
|
171
|
+
x_ = x.reshape(-1, self.d)
|
|
172
|
+
x_ = self._standardize_x(x_)
|
|
173
|
+
|
|
174
|
+
# Setup generator for reproducibility
|
|
175
|
+
g = None
|
|
176
|
+
if seed is not None:
|
|
177
|
+
g = torch.Generator(device=x.device)
|
|
178
|
+
g.manual_seed(int(seed))
|
|
179
|
+
|
|
180
|
+
# Sample latent variable u
|
|
181
|
+
if self.latent_dist == "uniform":
|
|
182
|
+
u = torch.rand(x_.shape[0], num_samples, self.q, generator=g, device=x.device)
|
|
183
|
+
else:
|
|
184
|
+
u = torch.randn(x_.shape[0], num_samples, self.q, generator=g, device=x.device)
|
|
185
|
+
|
|
186
|
+
# Get location parameter (mean of the kernel)
|
|
187
|
+
y_loc = self.forward(x_, u)
|
|
188
|
+
|
|
189
|
+
# Sample noise from Normal(0,1)
|
|
190
|
+
# We use the same generator 'g' if provided to ensure full reproducibility
|
|
191
|
+
noise = torch.randn(y_loc.shape, generator=g, device=x.device)
|
|
192
|
+
|
|
193
|
+
# Add kernel noise (bandwidth) *before* destandardization
|
|
194
|
+
# self.eps() is the learned bandwidth in standardized space
|
|
195
|
+
y = y_loc + self.eps() * noise
|
|
196
|
+
|
|
197
|
+
# Destandardize to get back to original data scale
|
|
198
|
+
y = self._destandardize_y(y)
|
|
199
|
+
|
|
200
|
+
return y if(len(x.shape)==2 or self.q>1) else y.flatten()
|
|
201
|
+
else:
|
|
202
|
+
device = self.eps().device
|
|
203
|
+
return self.sample_conditional(torch.tensor(x, device=device, dtype=torch.float32), num_samples = num_samples, seed = seed).cpu().numpy()
|
|
204
|
+
|
|
205
|
+
def fit(self, xs: torch.Tensor, ys: torch.Tensor, epochs: int = 1000, lr: float = 1e-3, m: int = 30, h0: float = 5e-2, val_split: float = 0.1):
|
|
206
|
+
if(isinstance(xs, torch.Tensor) and isinstance(ys, torch.Tensor)):
|
|
207
|
+
self._istraining = True
|
|
208
|
+
device = xs.device
|
|
209
|
+
|
|
210
|
+
# --- Validation Split Logic ---
|
|
211
|
+
if val_split > 0:
|
|
212
|
+
indices = torch.randperm(xs.shape[0])
|
|
213
|
+
v_size = int(xs.shape[0] * val_split)
|
|
214
|
+
train_idx, val_idx = indices[v_size:], indices[:v_size]
|
|
215
|
+
xt, yt = xs[train_idx], ys[train_idx]
|
|
216
|
+
xv, yv = xs[val_idx], ys[val_idx]
|
|
217
|
+
else:
|
|
218
|
+
xt, yt = xs, ys
|
|
219
|
+
xv, yv = None, None
|
|
220
|
+
|
|
221
|
+
best_val_loss = float('inf')
|
|
222
|
+
best_state = None
|
|
223
|
+
|
|
224
|
+
# Pre-training initialization (Using xt/yt for stats)
|
|
225
|
+
with torch.no_grad():
|
|
226
|
+
self.x_mean = xt.mean(dim=0)
|
|
227
|
+
self.x_std = xt.std(dim=0) + 1e-10
|
|
228
|
+
self.y_mean = yt.mean(dim=0)
|
|
229
|
+
self.y_std = yt.std(dim=0) + 1e-10
|
|
230
|
+
|
|
231
|
+
# CHANGE 3: Initialize eps in log space
|
|
232
|
+
eps0 = torch.tensor([h0] * self.q, device=device, dtype=torch.float32)
|
|
233
|
+
try:
|
|
234
|
+
self._eps_param.copy_(torch.log(eps0))
|
|
235
|
+
except: pass
|
|
236
|
+
|
|
237
|
+
opt = torch.optim.Adam([p for p in self.parameters() if p.requires_grad], lr=lr)
|
|
238
|
+
pbar = tqdm(range(epochs), desc="Training CPFN")
|
|
239
|
+
|
|
240
|
+
for epoch in pbar:
|
|
241
|
+
self.train()
|
|
242
|
+
opt.zero_grad()
|
|
243
|
+
loss = -self.logdensity(xt, yt, m, tilted=True).mean()
|
|
244
|
+
loss.backward()
|
|
245
|
+
opt.step()
|
|
246
|
+
|
|
247
|
+
# --- Validation Check ---
|
|
248
|
+
curr_loss = loss.item()
|
|
249
|
+
if xv is not None:
|
|
250
|
+
self.eval()
|
|
251
|
+
with torch.no_grad():
|
|
252
|
+
v_loss = -self.logdensity(xv, yv, m, tilted=False).mean().item()
|
|
253
|
+
if v_loss < best_val_loss:
|
|
254
|
+
best_val_loss = v_loss
|
|
255
|
+
best_state = copy.deepcopy(self.state_dict())
|
|
256
|
+
curr_loss = v_loss # Show validation loss in pbar
|
|
257
|
+
|
|
258
|
+
eps_vals = self.eps().detach().cpu().numpy()
|
|
259
|
+
eps_str = ", ".join([f"{e:.2e}" for e in eps_vals])
|
|
260
|
+
pbar.set_postfix({"loss": f"{curr_loss:.4e}", "bw": eps_str})
|
|
261
|
+
|
|
262
|
+
# Restore best weights if validation was used
|
|
263
|
+
if best_state is not None:
|
|
264
|
+
self.load_state_dict(best_state)
|
|
265
|
+
print(f"Restored best model with validation loss: {best_val_loss:.4e}")
|
|
266
|
+
|
|
267
|
+
self._istraining = False
|
|
268
|
+
else:
|
|
269
|
+
# Recursive call for numpy array inputs (updated signature)
|
|
270
|
+
device = self.eps().device
|
|
271
|
+
self.fit(torch.tensor(xs, device=device, dtype=torch.float32),
|
|
272
|
+
torch.tensor(ys, device=device, dtype=torch.float32),
|
|
273
|
+
epochs=epochs, lr=lr, m=m, h0=h0, val_split=val_split)
|
|
274
|
+
|
|
275
|
+
def freeze(self):
|
|
276
|
+
# Freeze all parameters (prevent gradient updates)
|
|
277
|
+
for p in self.parameters():
|
|
278
|
+
p.requires_grad_(False)
|
|
279
|
+
self.eval()
|
|
280
|
+
|
|
281
|
+
def unfreeze(self):
|
|
282
|
+
# Unfreeze all parameters (enable gradient updates)
|
|
283
|
+
for p in self.parameters():
|
|
284
|
+
p.requires_grad_(True)
|
|
285
|
+
self.train()
|
|
286
|
+
|
|
287
|
+
def to(self, *args, **kwargs):
|
|
288
|
+
super().to(*args, **kwargs)
|
|
@@ -0,0 +1,127 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: cpfn
|
|
3
|
+
Version: 1.0.0
|
|
4
|
+
Summary: Conditional Push-Forward Neural Network estimator
|
|
5
|
+
Author-email: tedescolor <tedescolor@gmail.com>
|
|
6
|
+
License: Copyright (c) 2018 The Python Packaging Authority
|
|
7
|
+
|
|
8
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
9
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
10
|
+
in the Software without restriction, including without limitation the rights
|
|
11
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
12
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
13
|
+
furnished to do so, subject to the following conditions:
|
|
14
|
+
|
|
15
|
+
The above copyright notice and this permission notice shall be included in all
|
|
16
|
+
copies or substantial portions of the Software.
|
|
17
|
+
|
|
18
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
19
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
20
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
21
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
22
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
23
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
24
|
+
SOFTWARE.
|
|
25
|
+
Project-URL: Repository, https://github.com/tedescolor/cpfn
|
|
26
|
+
Project-URL: Paper, https://arxiv.org/pdf/2511.14455
|
|
27
|
+
Requires-Python: >=3.8
|
|
28
|
+
Description-Content-Type: text/markdown
|
|
29
|
+
License-File: LICENCE
|
|
30
|
+
Requires-Dist: torch>=2.0
|
|
31
|
+
Requires-Dist: numpy>=1.20
|
|
32
|
+
Requires-Dist: tqdm>=4.60
|
|
33
|
+
Dynamic: license-file
|
|
34
|
+
|
|
35
|
+
# CPFN — Conditional Push-Forward Neural Network
|
|
36
|
+
|
|
37
|
+
Compact, importable implementation of a Conditional Push-Forward Neural Network (CPFN) estimator.
|
|
38
|
+
|
|
39
|
+
**Paper:** https://arxiv.org/pdf/2511.14455
|
|
40
|
+
|
|
41
|
+
## Goals
|
|
42
|
+
- Provide a lightweight `CPFN` class for estimating conditional generators.
|
|
43
|
+
- Expose a simple API for training and sampling.
|
|
44
|
+
|
|
45
|
+
## Install
|
|
46
|
+
|
|
47
|
+
### From PyPI
|
|
48
|
+
```bash
|
|
49
|
+
pip install cpfn
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
## Quick Usage
|
|
53
|
+
```python
|
|
54
|
+
import torch
|
|
55
|
+
import matplotlib.pyplot as plt
|
|
56
|
+
import numpy as np
|
|
57
|
+
import random
|
|
58
|
+
from cpfn import CPFN
|
|
59
|
+
|
|
60
|
+
# 0. Hardware Selection (CUDA for NVIDIA, MPS for Apple Silicon, or CPU) and Reproducibility
|
|
61
|
+
if torch.cuda.is_available():
|
|
62
|
+
device = torch.device("cuda")
|
|
63
|
+
elif torch.backends.mps.is_available():
|
|
64
|
+
device = torch.device("mps")
|
|
65
|
+
else:
|
|
66
|
+
device = torch.device("cpu")
|
|
67
|
+
|
|
68
|
+
SEED = 43
|
|
69
|
+
|
|
70
|
+
random.seed(SEED)
|
|
71
|
+
np.random.seed(SEED)
|
|
72
|
+
torch.manual_seed(SEED)
|
|
73
|
+
|
|
74
|
+
# 1. Define Synthetic Ground Truth (Branching Function)
|
|
75
|
+
def true_sample(x):
|
|
76
|
+
z = np.random.randn() # Gaussian noise
|
|
77
|
+
p = np.random.rand() # Uniform switch
|
|
78
|
+
|
|
79
|
+
# Conditional logic: creates two paths for x > 0.5
|
|
80
|
+
if x < 0.5 or p < 0.5:
|
|
81
|
+
return 10 * x * (x - 0.5) * (1.5 - x) + z * 0.3 * (1.3 - x)
|
|
82
|
+
else:
|
|
83
|
+
return 10 * x * (x - 0.5) * (0.8 - x) + z * 0.3 * (1.3 - x)
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
# 2. Generate Training Data
|
|
87
|
+
ntrain = 500
|
|
88
|
+
xs = np.random.rand(ntrain)
|
|
89
|
+
ys = np.array([true_sample(x) for x in xs])
|
|
90
|
+
|
|
91
|
+
# 3. Model Setup & Training
|
|
92
|
+
model = CPFN(d=1, q=1, r=20, width=50, hidden_layers=3, delta=1e-15)
|
|
93
|
+
model.to(device)
|
|
94
|
+
model.fit(xs, ys,
|
|
95
|
+
epochs=3000,
|
|
96
|
+
lr=1e-3,
|
|
97
|
+
m=30,
|
|
98
|
+
h0=5e-2)
|
|
99
|
+
|
|
100
|
+
model.freeze()
|
|
101
|
+
|
|
102
|
+
# 4. Inference: Generate 1 sample for every x in training set
|
|
103
|
+
# samples shape: (ntrain, 1, 1)
|
|
104
|
+
ys_gen = model.sample_conditional(xs, num_samples=1).flatten()
|
|
105
|
+
|
|
106
|
+
# 5. Visualization
|
|
107
|
+
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5), sharey=True)
|
|
108
|
+
|
|
109
|
+
ax1.scatter(xs, ys, alpha=0.6, s=15, label="Ground Truth")
|
|
110
|
+
ax1.set_title("Original Training Data")
|
|
111
|
+
ax1.set_xlabel("x")
|
|
112
|
+
ax1.set_ylabel("y")
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
## Tests
|
|
117
|
+
Run the included pytest smoke test:
|
|
118
|
+
```bash
|
|
119
|
+
pytest -q
|
|
120
|
+
```
|
|
121
|
+
|
|
122
|
+
## Development
|
|
123
|
+
- Source: `src/cpfn/`
|
|
124
|
+
- Tests: `tests/`
|
|
125
|
+
|
|
126
|
+
## License
|
|
127
|
+
See `LICENCE` in the repository root.
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
LICENCE
|
|
2
|
+
README.md
|
|
3
|
+
pyproject.toml
|
|
4
|
+
src/cpfn/__init__.py
|
|
5
|
+
src/cpfn/estimator.py
|
|
6
|
+
src/cpfn.egg-info/PKG-INFO
|
|
7
|
+
src/cpfn.egg-info/SOURCES.txt
|
|
8
|
+
src/cpfn.egg-info/dependency_links.txt
|
|
9
|
+
src/cpfn.egg-info/requires.txt
|
|
10
|
+
src/cpfn.egg-info/top_level.txt
|
|
11
|
+
tests/test_estimator.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
cpfn
|
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
import math
|
|
2
|
+
|
|
3
|
+
import torch
|
|
4
|
+
|
|
5
|
+
from cpfn import CPFN
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def test_cpfn_smoke_cpu():
|
|
9
|
+
device = torch.device("cpu")
|
|
10
|
+
torch.manual_seed(0)
|
|
11
|
+
|
|
12
|
+
n = 20
|
|
13
|
+
xs = torch.rand(n, 1, device=device)
|
|
14
|
+
ys = torch.sin(2 * math.pi * xs)
|
|
15
|
+
|
|
16
|
+
model = CPFN(d=1, q=1, r=4, width=16, hidden_layers=1, learn_eps=True)
|
|
17
|
+
model.to(device)
|
|
18
|
+
|
|
19
|
+
# short training to smoke-test functionality
|
|
20
|
+
model.fit(xs, ys, epochs=5, lr=1e-3, m=5, h0=1e-2)
|
|
21
|
+
|
|
22
|
+
out = model.sample_conditional(xs[:4], num_samples=3, seed=42)
|
|
23
|
+
assert out.shape == (4, 3, 1)
|
|
24
|
+
|
|
25
|
+
eps = model.eps().detach().cpu().numpy()
|
|
26
|
+
assert (eps > 0).all()
|