pulse-code 1.0.1__py3-none-any.whl

pulse/matcher.py

@@ -0,0 +1,180 @@
+"""Fuzzy keyword matching for finding the best CDC WONDER dataset or query."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+from pulse.catalog import BundledQuery, Catalog, Dataset
+
+# Synonym groups — any word in a group matches any other word in that group
+_SYNONYMS: list[set[str]] = [
+    {"opioid", "opioids", "opiate", "opiates"},
+    {"fentanyl", "synthetic opioid", "synthetic narcotic"},
+    {"heroin"},
+    {"drug", "drugs", "narcotic", "narcotics", "overdose"},
+    {"cocaine", "coke"},
+    {"meth", "methamphetamine", "amphetamine", "stimulant", "psychostimulant"},
+    {
+        "death",
+        "deaths",
+        "mortality",
+        "die",
+        "dying",
+        "fatal",
+        "fatality",
+        "fatalities",
+        "kill",
+    },
+    {"birth", "births", "natality", "born", "fertility", "fertilit"},
+    {"infant", "neonatal", "neonate", "baby", "babies", "newborn", "SIDS"},
+    {"maternal", "mother", "pregnancy", "pregnant", "childbirth", "obstetric"},
+    {"suicide", "suicidal", "self-harm", "self harm"},
+    {"cancer", "tumor", "neoplasm", "oncology"},
+    {"heart", "cardiac", "cardiovascular", "coronary"},
+    {"race", "racial", "ethnicity", "ethnic", "disparity", "disparities", "equity"},
+    {"vaccine", "vaccination", "immunization", "vaers", "adverse event"},
+    {"covid", "COVID", "COVID-19", "coronavirus", "pandemic"},
+    {"tick", "lyme", "lyme disease", "rocky mountain", "anaplasmosis", "ehrlichiosis"},
+    {
+        "environment",
+        "environmental",
+        "climate",
+        "temperature",
+        "heat",
+        "air quality",
+        "pollution",
+        "PM2.5",
+        "precipitation",
+        "rain",
+    },
+    {
+        "cancer",
+        "tumor",
+        "neoplasm",
+        "oncology",
+        "malignant",
+        "carcinoma",
+        "leukemia",
+        "lymphoma",
+        "breast cancer",
+        "lung cancer",
+        "colorectal",
+        "prostate",
+    },
+    {
+        "STD",
+        "STI",
+        "chlamydia",
+        "gonorrhea",
+        "syphilis",
+        "sexually transmitted",
+        "sexual health",
+    },
+    {"tuberculosis", "TB", "OTIS", "MDR-TB", "drug-resistant TB"},
+    {"fetal death", "stillbirth", "fetal", "stillborn"},
+    {"AIDS", "HIV", "epidemic"},
+    {"population", "census", "demographics", "projections", "estimates", "denominator"},
+    {"sex", "gender", "male", "female"},
+    {"state", "states", "county", "counties", "geographic", "geography", "region"},
+    {"year", "annual", "yearly", "trend", "trends"},
+    {"month", "monthly", "seasonal"},
+    {"age", "ages", "age group", "elderly", "children", "young"},
+    {"provisional", "recent", "latest", "current"},
+    {"historical", "history", "long-term"},
+]
+
+# Build reverse lookup: word -> synonym set index
+_WORD_TO_GROUP: dict[str, int] = {}
+for _i, _group in enumerate(_SYNONYMS):
+    for _word in _group:
+        _WORD_TO_GROUP[_word.lower()] = _i
+
+
+def _tokenize(text: str) -> set[str]:
+    """Extract tokens and expand synonyms into group IDs."""
+    raw_tokens = set(re.findall(r"[a-zA-Z0-9.]+", text.lower()))
+    expanded: set[str] = set(raw_tokens)
+    for token in raw_tokens:
+        idx = _WORD_TO_GROUP.get(token)
+        if idx is not None:
+            expanded.add(f"__group_{idx}__")
+    return expanded
+
+
+def _score(prompt_tokens: set[str], item_tokens: set[str]) -> float:
+    """Jaccard-like overlap score, weighted by token count."""
+    if not item_tokens:
+        return 0.0
+    intersection = prompt_tokens & item_tokens
+    union = prompt_tokens | item_tokens
+    return len(intersection) / len(union)
+
+
+@dataclass
+class DatasetMatch:
+    dataset: Dataset
+    score: float
+    reason: str
+
+
+@dataclass
+class QueryMatch:
+    query: BundledQuery
+    score: float
+
+
+def _year_coverage_bonus(prompt: str, ds: "Dataset") -> float:
+    """Boost score if the dataset's year range covers years mentioned in the prompt."""
+    import re as _re
+
+    years = [int(y) for y in _re.findall(r"\b(19\d{2}|20\d{2})\b", prompt)]
+    if not years:
+        return 0.0
+    year_end = ds.year_end or 2026
+    covered = sum(1 for y in years if ds.year_start <= y <= year_end)
+    return 0.15 * covered / len(years)
+
+
+def match_datasets(prompt: str, catalog: Catalog, top_n: int = 5) -> list[DatasetMatch]:
+    """Return top-N datasets matching the prompt, scored by keyword overlap."""
+    prompt_tokens = _tokenize(prompt)
+    results: list[DatasetMatch] = []
+
+    for ds in catalog.datasets():
+        item_tokens = _tokenize(
+            " ".join([ds.title, ds.topic, ds.subject, ds.year_range_label] + ds.tags)
+        )
+        score = _score(prompt_tokens, item_tokens) + _year_coverage_bonus(prompt, ds)
+        reason = _build_reason(prompt_tokens, item_tokens, ds)
+        results.append(DatasetMatch(dataset=ds, score=score, reason=reason))
+
+    results.sort(key=lambda x: x.score, reverse=True)
+    return results[:top_n]
+
+
+def match_queries(prompt: str, catalog: Catalog, top_n: int = 5) -> list[QueryMatch]:
+    """Return top-N bundled queries matching the prompt."""
+    prompt_tokens = _tokenize(prompt)
+    results: list[QueryMatch] = []
+
+    for q in catalog.queries():
+        item_tokens = _tokenize(
+            " ".join(
+                [q.description, q.topic, q.year_range, q.dataset_id]
+                + q.tags
+                + q.groupings
+            )
+        )
+        score = _score(prompt_tokens, item_tokens)
+        results.append(QueryMatch(query=q, score=score))
+
+    results.sort(key=lambda x: x.score, reverse=True)
+    return results[:top_n]
+
+
+def _build_reason(prompt_tokens: set[str], item_tokens: set[str], ds: Dataset) -> str:
+    matched = [t for t in ds.tags if _tokenize(t) & prompt_tokens]
+    if matched:
+        return f"Matched: {', '.join(matched[:4])}"
+    return f"{ds.year_range_label} | {ds.topic}"

pulse/queries/aids-cases-by-year-1981-1999-req.xml

@@ -0,0 +1,178 @@
+<?xml version="1.0" encoding="UTF-8"?><request-parameters>
+	<parameter>
+		<name>B_1</name>
+		<value>D12.V17-level1</value>
+	</parameter>
+	<parameter>
+		<name>B_2</name>
+		<value>*None*</value>
+	</parameter>
+	<parameter>
+		<name>B_3</name>
+		<value>*None*</value>
+	</parameter>
+	<parameter>
+		<name>B_4</name>
+		<value>*None*</value>
+	</parameter>
+	<parameter>
+		<name>B_5</name>
+		<value>*None*</value>
+	</parameter>
+	<parameter>
+		<name>F_D12.V14</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>F_D12.V17</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>I_D12.V14</name>
+		<value></value>
+	</parameter>
+	<parameter>
+		<name>I_D12.V17</name>
+		<value></value>
+	</parameter>
+	<parameter>
+		<name>M_1</name>
+		<value>D12.M1</value>
+	</parameter>
+	<parameter>
+		<name>O_export-format</name>
+		<value>xls</value>
+	</parameter>
+	<parameter>
+		<name>O_precision</name>
+		<value>0</value>
+	</parameter>
+	<parameter>
+		<name>O_timeout</name>
+		<value>600</value>
+	</parameter>
+	<parameter>
+		<name>O_javascript</name>
+		<value>on</value>
+	</parameter>
+	<parameter>
+		<name>O_title</name>
+		<value></value>
+	</parameter>
+	<parameter>
+		<name>O_oc-sect1-request</name>
+		<value>close</value>
+	</parameter>
+	<parameter>
+		<name>O_V14_fmode</name>
+		<value>freg</value>
+	</parameter>
+	<parameter>
+		<name>O_V17_fmode</name>
+		<value>freg</value>
+	</parameter>
+	<parameter>
+		<name>O_change_action-Send-Export Results</name>
+		<value>Export Results</value>
+	</parameter>
+	<parameter>
+		<name>O_show_totals</name>
+		<value>true</value>
+	</parameter>
+	<parameter>
+		<name>O_show_zeros</name>
+		<value>true</value>
+	</parameter>
+	<parameter>
+		<name>O_show_suppressed</name>
+		<value>true</value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V1</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V2</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V3</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V4</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V5</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V6</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V7</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V8</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V9</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V10</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V11</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V12</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V13</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V14</name>
+		<value></value>
+	</parameter>
+	<parameter>
+		<name>V_D12.V17</name>
+		<value></value>
+	</parameter>
+	<parameter>
+		<name>action-Send</name>
+		<value>Send</value>
+	</parameter>
+	<parameter>
+		<name>dataset_code</name>
+		<value>D12</value>
+	</parameter>
+	<parameter>
+		<name>dataset_label</name>
+		<value>AIDS Public Use (Vintage 1999)</value>
+	</parameter>
+	<parameter>
+		<name>dataset_vintage</name>
+		<value></value>
+	</parameter>
+	<parameter>
+		<name>finder-stage-D12.V14</name>
+		<value>codeset</value>
+	</parameter>
+	<parameter>
+		<name>finder-stage-D12.V17</name>
+		<value>codeset</value>
+	</parameter>
+	<parameter>
+		<name>stage</name>
+		<value>request</value>
+	</parameter>
+</request-parameters>

pulse/queries/births-by-year-1995-2002-req.xml

@@ -0,0 +1,226 @@
+<?xml version="1.0" encoding="UTF-8"?><request-parameters>
+	<parameter>
+		<name>B_1</name>
+		<value>D10.V20</value>
+	</parameter>
+	<parameter>
+		<name>B_2</name>
+		<value>*None*</value>
+	</parameter>
+	<parameter>
+		<name>B_3</name>
+		<value>*None*</value>
+	</parameter>
+	<parameter>
+		<name>B_4</name>
+		<value>*None*</value>
+	</parameter>
+	<parameter>
+		<name>B_5</name>
+		<value>*None*</value>
+	</parameter>
+	<parameter>
+		<name>F_D10.V21</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>F_D10.V22</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>F_D10.V37</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>I_D10.V21</name>
+		<value></value>
+	</parameter>
+	<parameter>
+		<name>I_D10.V22</name>
+		<value></value>
+	</parameter>
+	<parameter>
+		<name>I_D10.V37</name>
+		<value></value>
+	</parameter>
+	<parameter>
+		<name>M_1</name>
+		<value>D10.M1</value>
+	</parameter>
+	<parameter>
+		<name>O_export-format</name>
+		<value>xls</value>
+	</parameter>
+	<parameter>
+		<name>O_timeout</name>
+		<value>600</value>
+	</parameter>
+	<parameter>
+		<name>O_javascript</name>
+		<value>on</value>
+	</parameter>
+	<parameter>
+		<name>O_title</name>
+		<value></value>
+	</parameter>
+	<parameter>
+		<name>O_location</name>
+		<value>D10.V21</value>
+	</parameter>
+	<parameter>
+		<name>O_V21_fmode</name>
+		<value>freg</value>
+	</parameter>
+	<parameter>
+		<name>O_V22_fmode</name>
+		<value>freg</value>
+	</parameter>
+	<parameter>
+		<name>O_V37_fmode</name>
+		<value>freg</value>
+	</parameter>
+	<parameter>
+		<name>O_change_action-Send-Export Results</name>
+		<value>Export Results</value>
+	</parameter>
+	<parameter>
+		<name>O_show_totals</name>
+		<value>true</value>
+	</parameter>
+	<parameter>
+		<name>O_show_zeros</name>
+		<value>true</value>
+	</parameter>
+	<parameter>
+		<name>O_show_suppressed</name>
+		<value>true</value>
+	</parameter>
+	<parameter>
+		<name>O_precision</name>
+		<value>0</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V4</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V2</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V1</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V5</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V20</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V3</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V6</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V7</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V8</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V9</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V11</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V12</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V16</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V14</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V18</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V15</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V19</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V13</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V17</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V10</name>
+		<value>*All*</value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V21</name>
+		<value></value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V22</name>
+		<value></value>
+	</parameter>
+	<parameter>
+		<name>V_D10.V37</name>
+		<value></value>
+	</parameter>
+	<parameter>
+		<name>action-Send</name>
+		<value>Send</value>
+	</parameter>
+	<parameter>
+		<name>dataset_code</name>
+		<value>D10</value>
+	</parameter>
+	<parameter>
+		<name>dataset_label</name>
+		<value>Natality, 1995-2002 Request Form</value>
+	</parameter>
+	<parameter>
+		<name>dataset_vintage</name>
+		<value></value>
+	</parameter>
+	<parameter>
+		<name>finder-stage-D10.V21</name>
+		<value>codeset</value>
+	</parameter>
+	<parameter>
+		<name>finder-stage-D10.V22</name>
+		<value>codeset</value>
+	</parameter>
+	<parameter>
+		<name>finder-stage-D10.V37</name>
+		<value>codeset</value>
+	</parameter>
+	<parameter>
+		<name>stage</name>
+		<value>request</value>
+	</parameter>
+</request-parameters>