Minor code quality improvements

PolicyEngine · Oct 15, 2024 · 4365b1d · 4365b1d
1 parent 02df977
commit 4365b1d
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 40 deletions.
diff --git a/docs/validation.ipynb b/docs/validation.ipynb
diff --git a/policyengine_uk_data/datasets/frs/enhanced_frs.py b/policyengine_uk_data/datasets/frs/enhanced_frs.py
@@ -105,7 +105,7 @@ def dropout_weights(weights, p):
 
     start_loss = None
 
-    iterator = trange(5_000)
+    iterator = trange(1_000)
     for i in iterator:
         optimizer.zero_grad()
         weights_ = dropout_weights(weights, dropout_rate)

diff --git a/policyengine_uk_data/utils/imputations/capital_gains.py b/policyengine_uk_data/utils/imputations/capital_gains.py
@@ -1,5 +1,6 @@
 import pandas as pd
 import numpy as np
+from policyengine_core.data import Dataset
 
 # Fit a spline to each income band's percentiles
 try:
@@ -165,7 +166,7 @@ def stack_datasets(data_1, data_2):
     return joined_data
 
 
-def impute_cg_to_dataset(dataset):
+def impute_cg_to_dataset(dataset: Dataset):
     data = dataset.load_dataset()
     zero_weight_copy_1 = copy.deepcopy(data)
     zero_weight_copy_2 = copy.deepcopy(data)