Merge branch 'develop' into noise-function-docs

Author: NathanielBlairStahn

README.rst

@@ -19,8 +19,6 @@ or build it from source with
 
   ``> python setup.py install``
 
-This will make the ``pseudopeople`` library available to python and install a
-command-line executable called ``...`` that you can use to verify your
-installation with
-
-  ``> ... test``
+Documentation
+======================
+You can view documentation at https://pseudopeople.readthedocs.io/en/latest/

src/pseudopeople/entity_types.py

@@ -2,9 +2,11 @@
 from typing import Any, Callable, Dict
 
 import pandas as pd
+from loguru import logger
 from vivarium import ConfigTree
 from vivarium.framework.randomness import RandomnessStream
 
+from pseudopeople import schema_entities
 from pseudopeople.utilities import get_index_to_noise
 
 
@@ -53,6 +55,7 @@ class ColumnNoiseType:
     noise_function: Callable[[pd.Series, ConfigTree, RandomnessStream, Any], pd.Series]
     row_noise_level: float = 0.01
     token_noise_level: float = 0.1
+    noise_level_scaling_function: Callable[[str], float] = lambda x: 1.0
     additional_parameters: Dict[str, Any] = None
 
     def __call__(
@@ -62,18 +65,27 @@ def __call__(
         randomness_stream: RandomnessStream,
         additional_key: Any,
     ) -> pd.Series:
-        # TODO: this is a temporary hack to account for all string columns having been made categorical
-        #  We should record expected output dtype in the columns data structure
-        if column.dtype.name == "category":
-            column = column.astype(str)
-        else:
-            column = column.copy()
-        noise_level = configuration.row_noise_level
+        column = column.copy()
+        noise_level = configuration.row_noise_level * self.noise_level_scaling_function(
+            column.name
+        )
         to_noise_idx = get_index_to_noise(
             column, noise_level, randomness_stream, f"{self.name}_{additional_key}"
         )
+        if to_noise_idx.empty:
+            logger.debug(
+                f"No cells chosen to noise for noise function {self.name} on column {column.name}. "
+                "This is likely due to a combination of the configuration noise levels and the input data."
+            )
+            return column
         noised_data = self.noise_function(
             column.loc[to_noise_idx], configuration, randomness_stream, additional_key
         )
+
+        # Coerce noised column dtype back to original column's if it's changed
+        if noised_data.dtype.name != column.dtype.name:
+            noised_data = noised_data.astype(column.dtype)
+
         column.loc[to_noise_idx] = noised_data
+
         return column

src/pseudopeople/interface.py

@@ -1,7 +1,8 @@
 from pathlib import Path
-from typing import Union
+from typing import List, Union
 
 import pandas as pd
+import pyarrow.parquet as pq
 
 from pseudopeople.configuration import get_configuration
 from pseudopeople.constants import paths
@@ -14,6 +15,7 @@ def _generate_form(
     source: Union[Path, str, pd.DataFrame],
     seed: int,
     configuration: Union[Path, str, dict],
+    year_filter: dict,
 ) -> pd.DataFrame:
     """
     Helper for generating noised forms from clean data.
@@ -49,9 +51,11 @@ def _generate_form(
         data = source
     elif isinstance(source, Path):
         if source.suffix == ".hdf":
-            data = pd.read_hdf(source)
+            with pd.HDFStore(str(source), mode="r") as hdf_store:
+                data = hdf_store.select("data", where=year_filter["hdf"])
+            hdf_store.close()
         elif source.suffix == ".parquet":
-            data = pd.read_parquet(source)
+            data = pq.read_table(source, filters=year_filter["parquet"]).to_pandas()
         else:
             raise ValueError(
                 "Source path must either be a .hdf or a .parquet file. Provided "
@@ -64,7 +68,21 @@ def _generate_form(
             f"Source {source} must be either a pandas DataFrame or a path to a "
             "file containing a pandas DataFrame."
         )
-    return noise_form(form, data, configuration_tree, seed)
+
+    columns_to_keep = [c for c in form.columns]
+    # Coerce dtypes
+    for col in columns_to_keep:
+        if col.dtype_name != data[col.name].dtype.name:
+            data[col.name] = data[col.name].astype(col.dtype_name)
+    noised_form = noise_form(form, data, configuration_tree, seed)
+    noised_form = _extract_columns(columns_to_keep, noised_form)
+    return noised_form
+
+
+def _extract_columns(columns_to_keep, noised_form):
+    if columns_to_keep:
+        noised_form = noised_form[[c.name for c in columns_to_keep]]
+    return noised_form
 
 
 # TODO: add year as parameter to select the year of the decennial census to generate (MIC-3909)
@@ -75,93 +93,146 @@ def generate_decennial_census(
     source: Union[Path, str, pd.DataFrame] = None,
     seed: int = 0,
     configuration: Union[Path, str, dict] = None,
+    year: int = 2020,
 ) -> pd.DataFrame:
     """
     Generates noised decennial census data from un-noised data.
 
     :param source: A path to or pd.DataFrame of the un-noised source census data
     :param seed: An integer seed for randomness
     :param configuration: (optional) A path to a configuration YAML file or a dictionary to override the default configuration
+    :param year: The year from the data to noise
     :return: A pd.DataFrame of noised census data
     """
-    return _generate_form(FORMS.census, source, seed, configuration)
+    year_filter = {"hdf": None, "parquet": None}
+    if year:
+        year_filter["hdf"] = [f"{FORMS.census.date_column} == {year}."]
+        year_filter["parquet"] = [(FORMS.census.date_column, "==", year)]
+    return _generate_form(FORMS.census, source, seed, configuration, year_filter)
 
 
 def generate_american_communities_survey(
     source: Union[Path, str, pd.DataFrame] = None,
     seed: int = 0,
     configuration: Union[Path, str, dict] = None,
+    year: int = 2020,
 ) -> pd.DataFrame:
     """
     Generates noised American Communities Survey (ACS) data from un-noised data.
 
     :param source: A path to or pd.DataFrame of the un-noised source ACS data
     :param seed: An integer seed for randomness
     :param configuration: (optional) A path to a configuration YAML file or a dictionary to override the default configuration
+    :param year: The year from the data to noise
     :return: A pd.DataFrame of noised ACS data
     """
-    return _generate_form(FORMS.acs, source, seed, configuration)
+    year_filter = {"hdf": None, "parquet": None}
+    if year:
+        year_filter["hdf"] = [
+            f"{FORMS.acs.date_column} >= '{year}-01-01' and {FORMS.acs.date_column} <= '{year}-12-31'"
+        ]
+        year_filter["parquet"] = [
+            (FORMS.acs.date_column, ">=", pd.Timestamp(f"{year}-01-01")),
+            (FORMS.acs.date_column, "<=", pd.Timestamp(f"{year}-12-31")),
+        ]
+        seed = seed * 10_000 + year
+    return _generate_form(FORMS.acs, source, seed, configuration, year_filter)
 
 
 def generate_current_population_survey(
     source: Union[Path, str, pd.DataFrame] = None,
     seed: int = 0,
     configuration: Union[Path, str, dict] = None,
+    year: int = 2020,
 ) -> pd.DataFrame:
     """
     Generates noised Current Population Survey (CPS) data from un-noised data.
 
     :param source: A path to or pd.DataFrame of the un-noised source CPS data
     :param seed: An integer seed for randomness
     :param configuration: (optional) A path to a configuration YAML file or a dictionary to override the default configuration
+    :param year: The year from the data to noise
     :return: A pd.DataFrame of noised CPS data
     """
-    return _generate_form(FORMS.cps, source, seed, configuration)
+    year_filter = {"hdf": None, "parquet": None}
+    if year:
+        year_filter["hdf"] = [
+            f"{FORMS.cps.date_column} >= '{year}-01-01' and {FORMS.cps.date_column} <= '{year}-12-31'"
+        ]
+        year_filter["parquet"] = [
+            (FORMS.cps.date_column, ">=", pd.Timestamp(f"{year}-01-01")),
+            (FORMS.cps.date_column, "<=", pd.Timestamp(f"{year}-12-31")),
+        ]
+        seed = seed * 10_000 + year
+    return _generate_form(FORMS.cps, source, seed, configuration, year_filter)
 
 
 def generate_taxes_w2_and_1099(
     source: Union[Path, str, pd.DataFrame] = None,
     seed: int = 0,
     configuration: Union[Path, str, dict] = None,
+    year: int = 2020,
 ) -> pd.DataFrame:
     """
     Generates noised W2 and 1099 data from un-noised data.
 
     :param source: A path to or pd.DataFrame of the un-noised source W2 and 1099 data
     :param seed: An integer seed for randomness
     :param configuration: (optional) A path to a configuration YAML file or a dictionary to override the default configuration
+    :param year: The year from the data to noise
     :return: A pd.DataFrame of noised W2 and 1099 data
     """
-    return _generate_form(FORMS.tax_w2_1099, source, seed, configuration)
+    year_filter = {"hdf": None, "parquet": None}
+    if year:
+        year_filter["hdf"] = [f"{FORMS.tax_w2_1099.date_column} == {year}."]
+        year_filter["parquet"] = [(FORMS.tax_w2_1099.date_column, "==", year)]
+        seed = seed * 10_000 + year
+    return _generate_form(FORMS.tax_w2_1099, source, seed, configuration, year_filter)
 
 
 def generate_women_infants_and_children(
     source: Union[Path, str, pd.DataFrame] = None,
     seed: int = 0,
     configuration: Union[Path, str, dict] = None,
+    year: int = 2020,
 ) -> pd.DataFrame:
     """
     Generates noised Women Infants and Children (WIC) data from un-noised data.
 
     :param source: A path to or pd.DataFrame of the un-noised source WIC data
     :param seed: An integer seed for randomness
     :param configuration: (optional) A path to a configuration YAML file or a dictionary to override the default configuration
+    :param year: The year from the data to noise
     :return: A pd.DataFrame of noised WIC data
     """
-    return _generate_form(FORMS.wic, source, seed, configuration)
+    year_filter = {"hdf": None, "parquet": None}
+    if year:
+        year_filter["hdf"] = [f"{FORMS.wic.date_column} == {year}."]
+        year_filter["parquet"] = [(FORMS.wic.date_column, "==", year)]
+        seed = seed * 10_000 + year
+    return _generate_form(FORMS.wic, source, seed, configuration, year_filter)
 
 
 def generate_social_security(
     source: Union[Path, str, pd.DataFrame] = None,
     seed: int = 0,
     configuration: Union[Path, str, dict] = None,
+    year: int = 2020,
 ) -> pd.DataFrame:
     """
     Generates noised Social Security (SSA) data from un-noised data.
 
     :param source: A path to or pd.DataFrame of the un-noised source SSA data
     :param seed: An integer seed for randomness
     :param configuration: (optional) A path to a configuration YAML file or a dictionary to override the default configuration
+    :param year: The year up to which to noise from the data
     :return: A pd.DataFrame of noised SSA data
     """
-    return _generate_form(FORMS.ssa, source, seed, configuration)
+    year_filter = {"hdf": None, "parquet": None}
+    if year:
+        year_filter["hdf"] = [f"{FORMS.ssa.date_column} <= {year}."]
+        year_filter["parquet"] = [
+            (FORMS.ssa.date_column, "<=", pd.Timestamp(f"{year}-12-31"))
+        ]
+        seed = seed * 10_000 + year
+    return _generate_form(FORMS.ssa, source, seed, configuration, year_filter)

src/pseudopeople/noise_entities.py

@@ -1,6 +1,6 @@
 from typing import NamedTuple
 
-from pseudopeople import noise_functions
+from pseudopeople import noise_functions, utilities
 from pseudopeople.entity_types import ColumnNoiseType, RowNoiseType
 
 
@@ -24,6 +24,7 @@ class __NoiseTypes(NamedTuple):
     incorrect_selection: ColumnNoiseType = ColumnNoiseType(
         "incorrect_selection",
         noise_functions.generate_incorrect_selections,
+        noise_level_scaling_function=utilities.noise_scaling_incorrect_selection,
         token_noise_level=None,
     )
     # copy_from_within_household: ColumnNoiseType = ColumnNoiseType(