finished some parts to find possible to adjust in zone

Author: bobkatla

PopSynthesis/Methods/IPSF/SAA/operations/general.py

@@ -68,6 +68,7 @@ def init_syn_pop_saa(att:str, marginal_data: pd.DataFrame, pool: pd.DataFrame) -
 
 
 def adjust_atts_state_match_census(att: str, curr_syn_pop: Union[None, pd.DataFrame], census_data_by_att: pd.DataFrame, adjusted_atts: List[str], pool: pd.DataFrame) -> pd.DataFrame:
+    assert count_field in pool.columns
     updated_syn_pop = None
     if curr_syn_pop is None:
         updated_syn_pop = init_syn_pop_saa(att, census_data_by_att, pool).to_pandas()
@@ -76,7 +77,11 @@ def adjust_atts_state_match_census(att: str, curr_syn_pop: Union[None, pd.DataFr
         assert (states_diff_census.sum(axis=1) == 0).all()
         # With state diff we can now do adjustment for each zone, can parallel it?
         syn_pop_count = curr_syn_pop.value_counts().reset_index()
-        updated_syn_pop = zone_adjustment(curr_syn_pop)
+        for zid, zone_states_diff in states_diff_census.iterrows():
+            print(f"Processing zone {zid}")
+            sub_syn_pop = syn_pop_count[syn_pop_count[zone_field]==zid]
+            updated_syn_pop = zone_adjustment(att, sub_syn_pop, zone_states_diff, pool, adjusted_atts)
+            break
 
         # Will slowly convert to polars later
         # All the pool and synpop would be in the count format (with weights)

PopSynthesis/Methods/IPSF/SAA/operations/zone_adjustment.py

@@ -14,11 +14,43 @@
 """
 import pandas as pd
 import numpy as np
-from PopSynthesis.Methods.IPSF.const import zone_field
+from PopSynthesis.Methods.IPSF.const import count_field, zone_field
 from typing import List
+import itertools
+import random
 
-def zone_adjustment(curr_syn_count: pd.DataFrame, diff_census: pd.Series, pool: pd.DataFrame, adjusted_atts: List[str]) -> pd.DataFrame:
-    assert zone_field in curr_syn_count.columns
-    assert zone_field in pool.columns
 
-    print(diff_census)
+def find_combinations_with_prev_atts(filtered_targeted_samples: pd.DataFrame, prev_atts: List[str]):
+    sub_samples = filtered_targeted_samples[prev_atts + [count_field]]
+    count_samples = sub_samples.groupby(sub_samples.columns.difference([count_field]).to_list(), as_index=False)[count_field].sum()
+    return count_samples.set_index(count_samples.columns.difference([count_field]).to_list())
+
+def zone_adjustment(att: str, curr_syn_count: pd.DataFrame, diff_census: pd.Series, pool: pd.DataFrame, adjusted_atts: List[str]) -> pd.DataFrame:
+    assert count_field in curr_syn_count.columns
+    assert count_field in pool.columns
+    assert len(curr_syn_count[zone_field].unique()) == 1
+    curr_syn_count = curr_syn_count.drop(columns=[zone_field])
+
+    neg_states = diff_census[diff_census < 0].index.tolist()
+    pos_states = diff_census[diff_census > 0].index.tolist()
+    pairs_adjust = list(itertools.product(neg_states, pos_states))
+    random.shuffle(pairs_adjust)
+    
+    for neg_state, pos_state in pairs_adjust:
+        # Check neg state
+        filtered_syn_pop = curr_syn_count[curr_syn_count[att] == neg_state]
+        count_neg_in_syn = find_combinations_with_prev_atts(filtered_syn_pop, adjusted_atts)
+        neg_val = int(diff_census[neg_state])
+
+        # check pos state
+        filtered_pool = pool[pool[att] == pos_state]
+        count_pos_in_pool = find_combinations_with_prev_atts(filtered_pool, adjusted_atts)
+        pos_val = int(diff_census[pos_state])
+
+        n_adjust = min(abs(neg_val), pos_val) # the possible to adjust
+        possible_prev_comb = set(count_neg_in_syn.index) & set(count_pos_in_pool.index)
+        filtered_count_neg = count_neg_in_syn.loc[list(possible_prev_comb)]
+        filtered_count_pos = count_pos_in_pool.loc[list(possible_prev_comb)]
+        
+        
+        break