Coverage for src/crimepy/nibrs/munge.py: 86%
100 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-30 18:01 +0000
1"""Makes tables of NIBRS data ready for analysis in `pandas`.
3`prepare` is applied to every table that `crimepy.nibrs.read` and
4`crimepy.nibrs.load` return. It keeps every column and adds:
6* Every table: `state`, the postal abbreviation of the state.
7* Victims, offenders, and arrestees: `victim_age`, `offender_age`, and
8 `arrestee_age`, the age in years (0 for a baby under a year old, 99 for 99
9 or older, and missing if unknown). The age codes are kept as
10 `victim_age_code`, and so on.
11* Offenses, arrestees, and Group "B" arrests: `offense_category`,
12 `crime_against`, and `offense_group` of the offense (with the prefix
13 "arrest_" for arrests).
14* Agencies: `population`, the population of all of the agency's counties,
15 and `months_active`, the number of months in which the agency submitted a
16 report (of incidents or of no incidents).
18With `decode = True` (the default), codes are replaced with their labels
19(such as "Residence/Home" for location "20"), as categories. Codes that are
20not in a codebook are kept as they are. `decode` turns the codes of a table
21loaded with `decode = False` into labels later.
23Contents:
24 CODE_COLUMNS: the codebook of each column of codes.
25 add_ages: adds the age in years of each age code.
26 add_offense_details: adds the category, target, and group of offenses.
27 decode: replaces codes with their labels.
28 prepare: makes a table ready for analysis.
29 to_pandas: makes a `pyarrow` table into a `pandas` data frame.
31"""
33from __future__ import annotations
35from collections.abc import Callable, Mapping, Sequence
36from typing import Any, TypeVar, cast
38import amos
39import numpy as np
40import pandas as pd
41import pyarrow as pa
43from . import codes, layouts
45# The codebook of each column of codes, from the layouts.
46CODE_COLUMNS: dict[str, str] = {
47 f.name: f.codes
48 for layout in layouts.LAYOUTS.values() for f in layout
49 if f.kind == 'code' and f.codes is not None}
50# A table, or an `amos.Dataset` of one.
51_Table = TypeVar('_Table', pd.DataFrame, amos.Dataset)
52# Codebooks whose codes are kept, because another column has their meaning
53# (such as the age in years of each age code).
54_KEPT_CODEBOOKS: frozenset[str] = frozenset({'ages'})
55# Text columns with few values, which are made categories.
56_CATEGORY_TEXT: frozenset[str] = frozenset({
57 'state_abbreviation', 'city_submission', 'judicial_district',
58 'field_office', 'reporting_indicator'})
59_MONTHS: tuple[str, ...] = (
60 'jan', 'feb', 'mar', 'apr', 'may', 'jun', 'jul', 'aug', 'sep', 'oct',
61 'nov', 'dec')
62# The `pandas` type of each `pyarrow` type, so that whole numbers and
63# booleans with missing values keep their types.
64_PANDAS_TYPES: dict[pa.DataType, Any] = {
65 pa.int64(): pd.Int64Dtype(),
66 pa.bool_(): pd.BooleanDtype()}
69""" Public Functions """
72def to_pandas(table: pa.Table) -> pd.DataFrame:
73 """Returns a `pyarrow` table of NIBRS data as a `pandas` data frame.
75 Columns of codes become categories, whole numbers and booleans keep their
76 types (with `pd.NA` for missing values), and dates become `datetime64`.
78 Args:
79 table: a table made by `crimepy.nibrs.reader` or loaded from its
80 parquet files.
82 Returns:
83 The data frame.
85 """
86 columns = []
87 for name in table.column_names:
88 column = table.column(name)
89 if ((name in CODE_COLUMNS or name in _CATEGORY_TEXT)
90 and pa.types.is_string(column.type)):
91 column = column.dictionary_encode()
92 columns.append(column)
93 encoded = pa.table(columns, names = table.column_names)
94 return cast('pd.DataFrame', encoded.to_pandas(
95 types_mapper = _PANDAS_TYPES.get, date_as_object = False))
98def prepare(
99 data: pd.DataFrame,
100 table: str | None = None,
101 *,
102 decode: bool = True) -> pd.DataFrame:
103 """Makes a table of NIBRS data ready for analysis.
105 Args:
106 data: a table of NIBRS data (from `to_pandas`).
107 table: the name of the table (such as "victims"). Defaults to `None`,
108 in which case only the additions that every table gets are made.
109 decode: whether to replace codes with their labels. Defaults to
110 `True`.
112 Returns:
113 The table, with the added columns (see the module's documentation).
115 """
116 if 'state_code' in data.columns and 'state' not in data.columns: 116 ↛ 121line 116 didn't jump to line 121 because the condition on line 116 was always true
117 state = _lookup(data['state_code'], codes.STATE_ABBREVIATIONS.get)
118 data.insert(
119 _after(data, 'state_code'), 'state',
120 state.astype('category'))
121 data = add_ages(data)
122 data = add_offense_details(data)
123 if table == 'agencies':
124 data = _add_agency_totals(data)
125 return _decode(data) if decode else data
128def add_ages(data: pd.DataFrame) -> pd.DataFrame:
129 """Adds the age in years of each column of age codes.
131 A column "victim_age_code" gets a column "victim_age" after it, and so
132 on. Babies under a year old ("NN", "NB", and "BB") are 0, "99" (99 or
133 older) is 99, and an unknown age ("00") is missing.
135 Args:
136 data: a table with columns of age codes, whose names end in
137 "_age_code".
139 Returns:
140 The table, with the ages.
142 """
143 for column in [c for c in data.columns if str(c).endswith('_age_code')]:
144 name = str(column).removesuffix('_code')
145 if name in data.columns: 145 ↛ 146line 145 didn't jump to line 146 because the condition on line 145 was never true
146 continue
147 values = _lookup(data[column], _age_years)
148 data.insert(
149 _after(data, column), name, pd.to_numeric(values).astype('Int64'))
150 return data
153def add_offense_details(data: pd.DataFrame) -> pd.DataFrame:
154 """Adds the category, target, and group of each offense.
156 "offense_code" gets "offense_category", "crime_against", and
157 "offense_group" after it, and "arrest_offense" gets the same columns with
158 the prefix "arrest_". Only codes are looked up, so this is done before
159 the codes are decoded.
161 Args:
162 data: a table with an "offense_code" or "arrest_offense" column of
163 codes.
165 Returns:
166 The table, with the details.
168 """
169 for column, prefix in (('offense_code', ''), ('arrest_offense', 'arrest_')):
170 if (column not in data.columns
171 or f'{prefix}offense_category' in data.columns):
172 continue
173 offense = data[column]
174 position = _after(data, column)
175 for name, lookup in (
176 ('offense_category', codes.OFFENSE_CATEGORIES),
177 ('crime_against', codes.CRIMES_AGAINST),
178 ('offense_group', codes.OFFENSE_GROUPS)):
179 data.insert(
180 position, f'{prefix}{name}',
181 _lookup(offense, lookup.get).astype('category'))
182 position += 1
183 return data
186def decode(
187 data: _Table,
188 columns: Sequence[str] | None = None,
189 codebooks: Mapping[str, Mapping[str, str]] | None = None) -> _Table:
190 """Replaces codes with their labels.
192 Args:
193 data: a table with columns of codes, or an `amos.Dataset` of one
194 (whose data is decoded in place).
195 columns: the columns to decode. Defaults to `None`, which decodes
196 every column of codes in `CODE_COLUMNS` (except age codes, which
197 `add_ages` turns into years).
198 codebooks: codebooks to use instead of those in
199 `crimepy.nibrs.codes.CODEBOOKS`, by column. Defaults to `None`.
201 Returns:
202 The table (or dataset), with labels (as categories) instead of codes.
203 Codes that are not in a codebook are kept.
205 """
206 if isinstance(data, amos.Dataset): 206 ↛ 207line 206 didn't jump to line 207 because the condition on line 206 was never true
207 data.replace(decode(data.data, columns, codebooks))
208 return data
209 codebooks = codebooks or {}
210 if columns is None:
211 columns = [
212 c for c in data.columns
213 if c in CODE_COLUMNS and CODE_COLUMNS[c] not in _KEPT_CODEBOOKS]
214 for column in columns:
215 book = codebooks.get(column) or codes.CODEBOOKS[CODE_COLUMNS[column]]
216 data[column] = _relabel(data[column], book)
217 return data
220""" Private Functions """
223# `prepare` has a `decode` argument, so it calls `decode` by this name.
224_decode = decode
227def _add_agency_totals(data: pd.DataFrame) -> pd.DataFrame:
228 """Adds an agency's total population and number of active months.
230 Args:
231 data: the table of agencies.
233 Returns:
234 The table, with "population" and "months_active".
236 """
237 counties = [f'population_{n}' for n in range(1, 6)]
238 if (all(c in data.columns for c in counties) 238 ↛ 244line 238 didn't jump to line 244 because the condition on line 238 was always true
239 and 'population' not in data.columns):
240 position = _after(data, 'population_1') - 1
241 data.insert(
242 position, 'population',
243 data[counties].sum(axis = 1, min_count = 1).astype('Int64'))
244 flags = [
245 (f'zero_report_{m}', f'incident_report_{m}') for m in _MONTHS]
246 if all(a in data.columns and b in data.columns for a, b in flags): 246 ↛ 252line 246 didn't jump to line 252 because the condition on line 246 was always true
247 active = sum(
248 ((data[a].fillna(False) | data[b].fillna(False)).astype(int)
249 for a, b in flags),
250 start = pd.Series(0, index = data.index))
251 data['months_active'] = active.astype('Int64')
252 return data
255def _after(data: pd.DataFrame, column: str) -> int:
256 """Returns the position just after `column` in `data`.
258 Args:
259 data: a table.
260 column: a column of `data`.
262 Returns:
263 The position at which a column inserted after `column` goes.
265 """
266 return list(data.columns).index(column) + 1
269def _age_years(code: Any) -> float | None:
270 """Returns the age in years of an age code.
272 Args:
273 code: an age code (such as "25" or "NB"), or a missing value.
275 Returns:
276 The age in years, or `None` if it is unknown.
278 """
279 if not isinstance(code, str): 279 ↛ 280line 279 didn't jump to line 280 because the condition on line 279 was never true
280 return None
281 if code in codes.AGE_YEARS: 281 ↛ 282line 281 didn't jump to line 282 because the condition on line 281 was never true
282 return codes.AGE_YEARS[code]
283 return float(code) if code.isdigit() else None
286def _lookup(values: pd.Series, find: Callable[[Any], Any]) -> pd.Series:
287 """Returns what `find` returns for each value.
289 `find` is called once for each distinct value, so a column of categories
290 with millions of rows is looked up quickly.
292 Args:
293 values: the values to look up.
294 find: returns the result for a value (or `None`).
296 Returns:
297 The results, with the index of `values`.
299 """
300 if not isinstance(values.dtype, pd.CategoricalDtype): 300 ↛ 301line 300 didn't jump to line 301 because the condition on line 300 was never true
301 values = values.astype('category')
302 found = [find(c) for c in values.cat.categories] + [None]
303 codes_ = values.cat.codes.to_numpy()
304 return pd.Series(
305 np.asarray(found, dtype = object)[codes_], index = values.index)
308def _relabel(values: pd.Series, book: Mapping[str, str]) -> pd.Series:
309 """Returns codes replaced with their labels, as categories.
311 Args:
312 values: codes, as categories or text.
313 book: the label of each code.
315 Returns:
316 The labels, as categories in the order of `book`, followed by any
317 codes that are not in it.
319 """
320 if not isinstance(values.dtype, pd.CategoricalDtype): 320 ↛ 321line 320 didn't jump to line 321 because the condition on line 320 was never true
321 values = values.astype('category')
322 present = [str(c) for c in values.cat.categories]
323 labels = [book.get(c, c) for c in present]
324 if len(set(labels)) != len(labels): 324 ↛ 325line 324 didn't jump to line 325 because the condition on line 324 was never true
325 mapped = values.astype('object').map(lambda c: book.get(c, c))
326 return mapped.astype('category')
327 relabeled = values.cat.rename_categories(labels)
328 ordered = [v for v in book.values() if v in set(labels)]
329 ordered += [v for v in labels if v not in set(ordered)]
330 return relabeled.cat.reorder_categories(ordered)