diff --git a/changelog_entry.yaml b/changelog_entry.yaml index e69de29..ab1cb4c 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -0,0 +1,4 @@ +- bump: minor + changes: + added: + - Normalization parameter to handle multi-level geography calibration. diff --git a/src/microcalibrate/reweight.py b/src/microcalibrate/reweight.py index 3145a70..e56ae71 100644 --- a/src/microcalibrate/reweight.py +++ b/src/microcalibrate/reweight.py @@ -24,6 +24,7 @@ def reweight( epochs: Optional[int] = 2_000, noise_level: Optional[float] = 10.0, learning_rate: Optional[float] = 1e-3, + normalization_factor: Optional[torch.Tensor] = None, csv_path: Optional[str] = None, device: Optional[str] = None, ) -> tuple[np.ndarray, np.ndarray]: @@ -38,6 +39,9 @@ def reweight( epochs (int): Optional number of epochs for training. noise_level (float): Optional level of noise to add to the original weights. learning_rate (float): Optional learning rate for the optimizer. + normalization_factor (Optional[torch.Tensor]): Optional normalization factor for the loss (handles multi-level geographical calibration). + csv_path (Optional[str]): Optional path to save the performance metrics as a CSV file. + device (Optional[str]): Device to run the calibration on (e.g., 'cpu' or 'cuda'). If None, uses the default device. Returns: np.ndarray: Reweighted weights. @@ -103,7 +107,7 @@ def dropout_weights(weights: torch.Tensor, p: float) -> torch.Tensor: optimizer.zero_grad() weights_ = dropout_weights(weights, dropout_rate) estimate = estimate_function(torch.exp(weights_)) - l = loss(estimate, targets) + l = loss(estimate, targets, normalization_factor) close = pct_close(estimate, targets) if i % progress_update_interval == 0: diff --git a/src/microcalibrate/utils/metrics.py b/src/microcalibrate/utils/metrics.py index db4d975..da55edf 100644 --- a/src/microcalibrate/utils/metrics.py +++ b/src/microcalibrate/utils/metrics.py @@ -8,17 +8,23 @@ def loss( estimate: torch.Tensor, targets_array: torch.Tensor, + normalization_factor: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Calculate the loss based on the current weights and targets. Args: estimate (torch.Tensor): Current estimates in log space. targets_array (torch.Tensor): Array of target values. + normalization_factor (Optional[torch.Tensor]): Optional normalization factor for the loss (handles multi-level geographical calibration). Returns: torch.Tensor: Mean squared relative error between estimated and target values. """ rel_error = (((estimate - targets_array) + 1) / (targets_array + 1)) ** 2 + if normalization_factor is not None: + rel_error *= normalization_factor + if torch.isnan(rel_error).any(): + raise ValueError("Relative error contains NaNs") return rel_error.mean()