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 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,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,4 @@
1
+ from .estimator import CPFN
2
+
3
+ __all__ = ["CPFN"]
4
+
@@ -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,3 @@
1
+ torch>=2.0
2
+ numpy>=1.20
3
+ tqdm>=4.60
@@ -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()