add single year to data when extractinging dws

ihmeuw · Feb 20, 2024 · 43bc228 · 43bc228
1 parent 719d0c3
commit 43bc228
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/src/vivarium_inputs/core.py b/src/vivarium_inputs/core.py
@@ -196,9 +196,6 @@ def get_disability_weight(
                 validate=True,
                 get_all_years=get_all_years,
             )
-            # add year id with single year so normalization doesn't fill in all years
-            if not get_all_years:
-                data["year_id"] = utility_data.get_most_recent_year()
             data = utilities.normalize(data)
 
             cause = [c for c in causes if c.sequelae and entity in c.sequelae][0]

diff --git a/src/vivarium_inputs/extract.py b/src/vivarium_inputs/extract.py
@@ -196,6 +196,8 @@ def extract_disability_weight(
     data = disability_weights.loc[
         disability_weights.healthstate_id == entity.healthstate.gbd_id, :
     ]
+    if not get_all_years:
+        data["year_id"] = get_most_recent_year()
     return data