Add state space

Author: pockerman

src/spaces/state_space.py

@@ -0,0 +1,64 @@
+"""
+Discretized state space
+"""
+
+from typing import TypeVar, List
+from gym.spaces.discrete import Discrete
+
+from src.exceptions.exceptions import Error
+
+ActionStatus = TypeVar("ActionStatus")
+Env = TypeVar("Env")
+
+
+class State(object):
+    """
+    Describes an environment state
+    """
+    def __init__(self, column_name: str, state_id: int):
+        self.column_name: str = column_name
+        self.state_id: int = state_id
+        self.history: List[ActionStatus] = []
+
+    @property
+    def key(self) -> tuple:
+        return self.column_name, self.state_id
+
+
+class StateSpace(Discrete):
+    """
+    The State space is accumulates the discrete states
+    """
+
+    def __init__(self):
+        super(StateSpace, self).__init__(n=0)
+        self.states = {}
+
+    def init_from_environment(self, env: Env):
+        """
+        Initialize from environment
+        :param env:
+        :return:
+        """
+        names = env.feature_names
+        for col_name in names:
+
+            if col_name in self.states:
+                raise ValueError("Column {0} already exists".format(col_name))
+
+            self.states[col_name] = State(column_name=col_name, state_id=len(self.states))
+
+        # set the number of discrete states
+        self.n = len(self.states)
+
+    def add_state(self, state: State):
+        if state.column_name in self.states:
+            raise ValueError("Column {0} already exists".format(state.column_name))
+
+        self.states[state.column_name] = state
+
+    def update_state(self, state_name, status: ActionStatus):
+        self.states[state_name].history.append(status)
+
+    def __len__(self):
+        return len(self.states)

src/tests/test_space_state.py

@@ -0,0 +1,77 @@
+import unittest
+
+import unittest
+from pathlib import Path
+
+import pytest
+
+from src.spaces.environment import Environment
+from src.spaces.action_space import ActionSpace
+from src.spaces.actions import ActionSuppress, ActionGeneralize
+from src.exceptions.exceptions import Error
+from src.utils.serial_hierarchy import SerialHierarchy
+from src.utils.string_distance_calculator import DistanceType
+from src.datasets.dataset_wrapper import PandasDSWrapper
+from src.spaces.state_space import StateSpace, State
+
+class TestStateSpace(unittest.TestCase):
+
+    def setUp(self) -> None:
+        """
+        Setup the PandasDSWrapper to be used in the tests
+        :return: None
+        """
+
+        # read the data
+        filename = Path("../../data/mocksubjects.csv")
+
+        cols_types = {"gender": str, "ethnicity": str, "education": int,
+                      "salary": int, "diagnosis": int, "preventative_treatment": str,
+                      "mutation_status": int, }
+
+        self.ds = PandasDSWrapper(columns=cols_types)
+        self.ds.read(filename=filename, **{"features_drop_names": ["NHSno", "given_name", "surname", "dob"],
+                                           "names": ["NHSno", "given_name", "surname", "gender",
+                                                     "dob", "ethnicity", "education", "salary",
+                                                     "mutation_status", "preventative_treatment", "diagnosis"],
+                                           "drop_na": True,
+                                           "change_col_vals": {"diagnosis": [('N', 0)]}})
+
+    def test_creation(self):
+
+        action_space = ActionSpace(n=3)
+
+        generalization_table = {"Mixed White/Asian": SerialHierarchy(values=["Mixed", ]),
+                                "Chinese": SerialHierarchy(values=["Asian", ]),
+                                "Indian": SerialHierarchy(values=["Asian", ]),
+                                "Mixed White/Black African": SerialHierarchy(values=["Mixed", ]),
+                                "Black African": SerialHierarchy(values=["Black", ]),
+                                "Asian other": SerialHierarchy(values=["Asian", ]),
+                                "Black other": SerialHierarchy(values=["Black", ]),
+                                "Mixed White/Black Caribbean": SerialHierarchy(values=["Mixed", ]),
+                                "Mixed other": SerialHierarchy(values=["Mixed", ]),
+                                "Arab": SerialHierarchy(values=["Asian", ]),
+                                "White Irish": SerialHierarchy(values=["White", ]),
+                                "Not stated": SerialHierarchy(values=["Not stated"]),
+                                "White Gypsy/Traveller": SerialHierarchy(values=["White", ]),
+                                "White British": SerialHierarchy(values=["White", ]),
+                                "Bangladeshi": SerialHierarchy(values=["Asian", ]),
+                                "White other": SerialHierarchy(values=["White", ]),
+                                "Black Caribbean": SerialHierarchy(values=["Black", ]),
+                                "Pakistani": SerialHierarchy(values=["Asian", ])}
+
+        action_space.add(ActionGeneralize(column_name="ethnicity", generalization_table=generalization_table))
+
+        # create the environment from the given dataset
+        env = Environment(data_set=self.ds, action_space=action_space, gamma=0.99, start_column="gender")
+
+        state_space = StateSpace()
+        state_space.init_from_environment(env=env)
+
+        print(state_space.states.keys())
+
+        self.assertEqual(env.n_features, state_space.n)
+
+
+if __name__ == '__main__':
+    unittest.main()

src/utils/reward_manager.py

@@ -0,0 +1,14 @@
+"""
+Various utilities to handle reward assignment
+"""
+
+
+class RewardManager(object):
+    """
+    Helper class to assign rewards
+    """
+    def __init__(self) -> None:
+        pass
+
+    def get_state_reward(self, *options) -> float:
+        return 0.0