microcalibrate 0.1.0__py3-none-any.whl

microcalibrate/__init__.py

@@ -0,0 +1 @@
+from .calibration import Calibration

microcalibrate/calibration.py

@@ -0,0 +1,27 @@
+import pandas as pd
+import numpy as np
+
+
+class Calibration:
+    def __init__(
+        self,
+        data: pd.DataFrame,
+        weights: np.ndarray,
+        targets: np.ndarray,
+    ):
+        self.data = data
+        self.weights = weights
+        self.targets = targets
+
+    def calibrate(self):
+        from .reweight import reweight
+
+        new_weights, subsample = reweight(
+            original_weights=self.weights,
+            loss_matrix=self.data,
+            targets_array=self.targets,
+            epochs=32,
+        )
+
+        self.data = self.data.loc[subsample]
+        self.weights = new_weights

microcalibrate/reweight.py

@@ -0,0 +1,109 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+
+# Add device variable to use gpu (incl mps) if available
+device = torch.device(
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps" if torch.backends.mps.is_available() else "cpu"
+)
+
+
+def reweight(
+    original_weights,
+    loss_matrix,
+    targets_array,
+    dropout_rate=0.1,
+    epochs=2_000,
+    noise_level=10.0,
+    subsample_every=50,
+    learning_rate=1e-3,
+):
+    target_names = np.array(loss_matrix.columns)
+    loss_matrix = torch.tensor(
+        loss_matrix.values, dtype=torch.float32, device=device
+    )
+    original_indices = np.arange(loss_matrix.shape[0])
+
+    targets_array = torch.tensor(
+        targets_array, dtype=torch.float32, device=device
+    )
+    random_noise = np.random.random(original_weights.shape) * noise_level
+    weights = torch.tensor(
+        np.log(original_weights + random_noise),
+        requires_grad=True,
+        dtype=torch.float32,
+        device=device,
+    )
+
+    def loss(weights):
+        estimate = weights @ loss_matrix
+        rel_error = (
+            ((estimate - targets_array) + 1) / (targets_array + 1)
+        ) ** 2
+        return rel_error.mean()
+
+    def dropout_weights(weights, p):
+        if p == 0:
+            return weights
+        total_weight = weights.sum()
+        mask = torch.rand_like(weights) < p
+        masked_weights = weights.clone()
+        masked_weights[mask] = 0
+        masked_weights = masked_weights / masked_weights.sum() * total_weight
+        return masked_weights
+
+    optimizer = torch.optim.Adam([weights], lr=learning_rate)
+
+    iterator = tqdm(range(epochs), desc="Reweighting progress", unit="epoch")
+
+    for i in iterator:
+        optimizer.zero_grad()
+        running_loss = None
+        for j in range(2):
+            weights_ = dropout_weights(weights, dropout_rate)
+            l = loss(torch.exp(weights_))
+            if running_loss is None:
+                running_loss = l
+            else:
+                running_loss += l
+        l = running_loss / 2
+
+        if i % 10 == 0:
+            iterator.set_postfix(
+                {
+                    "loss": l.item(),
+                    "count_observations": loss_matrix.shape[0],
+                    "weights_mean": torch.exp(weights).mean().item(),
+                    "weights_std": torch.exp(weights).std().item(),
+                    "weights_min": torch.exp(weights).min().item(),
+                }
+            )
+
+        l.backward()
+        optimizer.step()
+
+        if subsample_every > 0 and i % subsample_every == 0 and i > 0:
+            weight_values = np.exp(weights.detach().cpu().numpy())
+
+            k = 100
+            # indices = indices of weights with values < 1
+            indices = np.where(weight_values >= k)[0]
+            loss_matrix = loss_matrix[indices, :]
+            weights = weights[indices]
+
+            loss_matrix = torch.tensor(
+                loss_matrix.detach().cpu(), dtype=torch.float32, device=device
+            )
+            weights = torch.tensor(
+                weights.detach().cpu(),
+                dtype=torch.float32,
+                device=device,
+                requires_grad=True,
+            )
+
+            original_indices = original_indices[indices]
+            optimizer = torch.optim.Adam([weights], lr=learning_rate)
+
+    return torch.exp(weights).detach().cpu().numpy(), original_indices

microcalibrate/target_correlation.py

@@ -0,0 +1,58 @@
+import numpy as np
+import torch
+import pandas as pd
+
+device = torch.device(
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps" if torch.backends.mps.is_available() else "cpu"
+)
+
+
+def analyse_targets(
+    weights,
+    loss_matrix,
+    targets_array,
+):
+    target_names = np.array(loss_matrix.columns)
+    loss_matrix = torch.tensor(
+        loss_matrix.values, dtype=torch.float32, device=device
+    )
+    targets_array = torch.tensor(
+        targets_array, dtype=torch.float32, device=device
+    )
+    weights = torch.tensor(
+        np.log(weights), requires_grad=True, dtype=torch.float32, device=device
+    )
+
+    # Compute estimates
+    estimate = torch.exp(weights) @ loss_matrix
+
+    # Compute relative errors for each target
+    rel_errors = ((estimate - targets_array) / targets_array) ** 2
+
+    # Use jacobian to compute gradients for all targets at once
+    gradients = torch.autograd.functional.jacobian(
+        func=lambda w: (
+            (torch.exp(w) @ loss_matrix - targets_array) / targets_array
+        )
+        ** 2,
+        inputs=weights,
+    )
+
+    # gradients is now [n_targets, n_weights]
+    gradients_numpy = gradients.cpu().numpy()
+    gradients_df = pd.DataFrame(gradients_numpy, index=target_names)
+
+    # Calculate correlation matrix
+    target_correlation = gradients_df.T.corr().mean()
+
+    df = pd.DataFrame(
+        {
+            "target": target_names,
+            "average_correlation": target_correlation,
+            "relative_error": rel_errors.detach().cpu().numpy() ** 0.5,
+        }
+    )
+
+    return df

microcalibrate-0.1.0.dist-info/METADATA

@@ -0,0 +1,16 @@
+Metadata-Version: 2.4
+Name: microcalibrate
+Version: 0.1.0
+Summary: Add your description here
+Author-email: Nikhil Woodruff <nikhil.woodruff@outlook.com>
+Requires-Python: >=3.11
+Requires-Dist: torch>=2.7.0
+Provides-Extra: dev
+Requires-Dist: black>=23.0.0; extra == 'dev'
+Requires-Dist: build<2.0.0,>=1.0.0; extra == 'dev'
+Requires-Dist: flake8<7.0.0,>=6.0.0; extra == 'dev'
+Requires-Dist: isort<6.0.0,>=5.9.0; extra == 'dev'
+Requires-Dist: linecheck<0.2.0,>=0.1.0; extra == 'dev'
+Requires-Dist: mypy<2.0.0,>=1.0.0; extra == 'dev'
+Requires-Dist: pytest-cov<7.0.0,>=6.0.0; extra == 'dev'
+Requires-Dist: pytest<9.0.0,>=8.0.0; extra == 'dev'

microcalibrate-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+microcalibrate/__init__.py,sha256=NP_5n7HBFIF1U5goXpLo3loPbmAtWudUzK_tRYaNYf0,37
+microcalibrate/calibration.py,sha256=hOMwkzvIX8V-SoF_uGt9vJTqPDoiBYgT_WjNiqGtnp8,613
+microcalibrate/reweight.py,sha256=QlK21Wy3GC5h-iSLs1ap_wQTTlALq_XpG9_7ku0LY9U,3380
+microcalibrate/target_correlation.py,sha256=NxhnzCMyjUMEUEUBUdazMEVwZF3W3SGpxEt32BKBSok,1558
+microcalibrate-0.1.0.dist-info/METADATA,sha256=lzIqNK4fvn-gBwRwo4CNB2K4kfIZlQvcZuY1RzZQ2Fc,637
+microcalibrate-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+microcalibrate-0.1.0.dist-info/entry_points.txt,sha256=rA66BJoIufeupyLVJq266BJ-jqy7YlQKovKTwY0Z__A,71
+microcalibrate-0.1.0.dist-info/RECORD,,

microcalibrate-0.1.0.dist-info/WHEEL

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

microcalibrate-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+policyengine-calibrate = policyengine_calibrate:main