yield_pyro.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "from aws_helper_functions import aws_helper_functions\n",
    "from sklearn.neighbors import BallTree\n",
    "from sklearn.model_selection import train_test_split\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "from pyro.nn import PyroModule\n",
    "import pyro.distributions as dist\n",
    "import pyro\n",
    "\n",
    "assert issubclass(PyroModule[nn.Linear], nn.Linear)\n",
    "assert issubclass(PyroModule[nn.Linear], PyroModule)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _retrieve_nearest_census_tract_numbers(df, local_mode=''):\n",
    "    # retrive census tract information from redshift\n",
    "    df_census = aws_helper_functions.read_from_redshift('SELECT * FROM raw_data_science.raw_commute_census_tracts_lat_long', local_mode=local_mode)\n",
    "    # Create a BallTree for the census tract latitudes and longitudes\n",
    "    tree = BallTree(df_census[['lat_orig', 'long_orig']].values, leaf_size=40)\n",
    "    \n",
    "    #drop rows in df where students_home__latitude__s or students_home__longitude__s is null\n",
    "    #df = df.dropna(subset=['latitude', 'longitude'])\n",
    "    df['latitude'] = np.where(df['latitude'].isna()==True, 40.776676, df['latitude'])\n",
    "    df['longitude'] = np.where(df['longitude'].isna()==True, -73.971321, df['longitude'])\n",
    "    \n",
    "    distances, indices = tree.query(df[['latitude', 'longitude']].values, k=1)\n",
    "    df.loc[:, 'boro_int'] = df_census.loc[indices.flatten(), 'boro_int'].values.copy()\n",
    "    df.loc[:, 'census_tract_int'] = df_census.loc[indices.flatten(), 'census_tract_int'].values.copy()\n",
    "    return df\n",
    "\n",
    "def _get_commute_time(df, local_mode=''):\n",
    "    #retrieve school-census tract commute times from redshift\n",
    "    df_commutes = aws_helper_functions.read_from_redshift('SELECT * FROM raw_data_science.raw_commute_census_tracts_to_schools', local_mode=local_mode)\n",
    "    \n",
    "    schools = _get_schools()\n",
    "    df_commutes = _replace_with_keys(df_commutes, 'school', schools)\n",
    "    #set df_commutes.time_walking_min, time_transit_min, and time_driving_min to numeric\n",
    "    df_commutes['time_walking_min'] = pd.to_numeric(df_commutes['time_walking_min'], errors='coerce')\n",
    "    df_commutes['time_transit_min'] = pd.to_numeric(df_commutes['time_transit_min'], errors='coerce')\n",
    "    df_commutes['time_driving_min'] = pd.to_numeric(df_commutes['time_driving_min'], errors='coerce')\n",
    "    \n",
    "    df_commutes['commute_time'] = df_commutes[['time_walking_min', 'time_transit_min']].min(axis=1)\n",
    "    df_commutes = df_commutes[['boro_int', 'census_tract_int', 'school', 'commute_time']]\n",
    "    \n",
    "    df = df.merge(df_commutes, how='left', on=['boro_int', 'census_tract_int', 'school'])\n",
    "    df['commute_time'] = np.where(df['commute_time']>120,30,df['commute_time'])\n",
    "    return df\n",
    "\n",
    "def _get_schools():\n",
    "    schools = {\n",
    "        'SA Bed-Stuy 2': 'BED-STUY2',\n",
    "        'SA Bed-Stuy': 'BED-STUY2',\n",
    "        'SA Bed-Stuy Middle School': 'BED-STUY_MIDDLE_SCHOOL',\n",
    "        'SA Bensonhurst': 'BENSONHURST',\n",
    "        'SA Bergen Beach':'BERGEN_BEACH',\n",
    "        'SA Bronx 1 Middle School': 'BRONX1',\n",
    "        'SA Bronx 1': 'BRONX1',\n",
    "        'SA Bronx Middle School': 'BRONX_MIDDLE_SCHOOL',\n",
    "        'SA Bronx 2': 'BRONX2',\n",
    "        'SA Bronx 2 Middle School': 'BRONX2_MIDDLE_SCHOOL',\n",
    "        'SA Bronx 3': 'BRONX3',\n",
    "        'SA Bronx 4': 'BRONX4',\n",
    "        'SA Bronx 5': 'BRONX5',\n",
    "        'SA Bronx 5 Upper': 'BRONX5',\n",
    "        'SA Bronx 5 Lower': 'BRONX5',\n",
    "        'SA Bushwick': 'BUSHWICK',\n",
    "        'SA Cobble Hill': 'COBBLE_HILL',\n",
    "        'SA Crown Heights': 'CROWN_HEIGHTS',\n",
    "        'SA Ditmas Park Middle School': 'DITMAS_PARK_MIDDLE_SCHOOL',\n",
    "        'SA East Flatbush Middle School': 'EAST_FLATBUSH_MIDDLE_SCHOOL',\n",
    "        'SA Far Rockaway': 'FAR_ROCKAWAY',\n",
    "        'SA Far Rockaway Middle School': 'FAR_ROCKAWAY_MIDDLE_SCHOOL',\n",
    "        'SA Flatbush': 'FLATBUSH',\n",
    "        'SA Hamilton Heights Middle School': 'HARLEM6',\n",
    "        'SA Harlem 1': 'HARLEM1',\n",
    "        'SA Harlem 2': 'HARLEM2',\n",
    "        'SA Harlem 3': 'HARLEM3',\n",
    "        'SA Harlem 4': 'HARLEM4',\n",
    "        'SA Harlem 5': 'HARLEM5',\n",
    "        'SA Harlem 6': 'HARLEM6',\n",
    "        'SA Harlem East': 'HARLEM_EAST',\n",
    "        'SA Harlem East Middle School': 'HARLEM_EAST',\n",
    "        'SA Harlem North Central': 'HARLEM_NORTH_CENTRAL',\n",
    "        'SA Harlem North Central Middle School': 'HARLEM_NORTH_CENTRAL',\n",
    "        'SA Harlem West': 'HARLEM_WEST',\n",
    "        'SA Harlem West Middle School': 'HARLEM_WEST',\n",
    "        'SA Harlem North West': 'HARLEM_NORTH_WEST',\n",
    "        'SA Harlem North West Middle School': 'HARLEM_NORTH_WEST',\n",
    "        'SA Hells Kitchen': 'HELLS_KITCHEN',\n",
    "        'SA Hell\\'s Kitchen': 'HELLS_KITCHEN',\n",
    "        'SA High School of the Liberal Arts - Manhattan': 'HIGH_SCHOOL_OF_THE_LIBERAL_ARTS_-_MANHATTAN',\n",
    "        'SA High School of the Liberal Arts-Manhattan': 'HIGH_SCHOOL_OF_THE_LIBERAL_ARTS_-_MANHATTAN',\n",
    "        'SA High School of the Liberal Arts - Harlem': 'HIGH_SCHOOL_OF_THE_LIBERAL_ARTS_-_HARLEM',\n",
    "        'SA High School of the Liberal Arts-Harlem': 'HIGH_SCHOOL_OF_THE_LIBERAL_ARTS_-_HARLEM',\n",
    "        'SA High School of the Liberal Arts - Brooklyn': 'HIGH_SCHOOL_OF_THE_LIBERAL_ARTS_-_BROOKLYN',\n",
    "        'SA High School of the Liberal Arts-Brooklyn': 'HIGH_SCHOOL_OF_THE_LIBERAL_ARTS_-_BROOKLYN',\n",
    "        'SA Hudson Yards': 'HUDSON_YARDS',\n",
    "        'SA Hudson Yards Middle School': 'HUDSON_YARDS_MIDDLE_SCHOOL',\n",
    "        'SA Kingsbridge Heights': 'KINGSBRIDGE_HEIGHTS',\n",
    "        'SA Lafayette Middle School': 'LAFAYETTE_MIDDLE_SCHOOL',\n",
    "        'SA Midtown West Middle School': 'MIDTOWN_WEST',\n",
    "        'SA Myrtle Middle School': 'MYRTLE_MIDDLE_SCHOOL',\n",
    "        'SA Norwood': 'NORWOOD',\n",
    "        'SA Ozone Park Middle School': 'OZONE_PARK_MIDDLE_SCHOOL',\n",
    "        'SA Prospect Heights': 'PROSPECT_HEIGHTS',\n",
    "        'SA Queens Village': 'QUEENS_VILLAGE',\n",
    "        'SA Rosedale': 'ROSEDALE',\n",
    "        'SA Rockaway Park Middle School': 'ROCKAWAY_PARK_MIDDLE_SCHOOL',\n",
    "        'SA South Jamaica': 'SOUTH_JAMAICA',\n",
    "        'SA Sheepshead Bay': 'SHEEPSHEAD_BAY',\n",
    "        'SA Springfield Gardens Middle School': 'SPRINGFIELD_GARDENS',\n",
    "        'SA Springfield Gardens MS': 'SPRINGFIELD_GARDENS',\n",
    "        'SA Springfield Gardens': 'SPRINGFIELD_GARDENS',\n",
    "        'SA Union Square': 'UNION_SQUARE',\n",
    "        'SA Upper West': 'UPPER_WEST',\n",
    "        'SA Washington Heights': 'WASHINGTON_HEIGHTS',\n",
    "        'SA Williamsburg': 'WILLIAMSBURG',\n",
    "        }\n",
    "    return schools\n",
    "\n",
    "def _replace_with_keys(df, column, dictionary):\n",
    "    new_df = pd.DataFrame()\n",
    "    for key, value in dictionary.items():\n",
    "        temp_df = df[df[column] == value].copy()\n",
    "        temp_df[column] = key\n",
    "        new_df = pd.concat([new_df, temp_df])\n",
    "    return new_df\n",
    "\n",
    "def _get_table_data(basetable=''):\n",
    "    raw_basetable = pd.read_csv(basetable)\n",
    "    df = raw_basetable\n",
    "    df = df[~df['grade'].isin(['PK', 'Not In School'])]\n",
    "    df = df[~df['grade'].isnull()]\n",
    "    df['scholar_grade'] = np.where(df['grade']=='K','0',df['grade']).astype(int)\n",
    "    df['latitude'] = df['students_home__latitude__s']\n",
    "    df['longitude'] = df['students_home__longitude__s']\n",
    "    df['school'] = df['accepted_school']\n",
    "    df = _get_commute_time(_retrieve_nearest_census_tract_numbers(df,local_mode=True),local_mode=True)\n",
    "    df['intercept'] = 1\n",
    "    df['es_school'] = df['scholar_grade'].isin(np.arange(0,5)).astype(int)\n",
    "    df['ms_school'] = df['scholar_grade'].isin(np.arange(5,6)).astype(int)\n",
    "    df['log_commute'] = np.log(df['commute_time'])\n",
    "    df['log_commute_square'] = df['log_commute'] ** 2\n",
    "    df['log_commute_third'] = df['log_commute'] ** 3\n",
    "    df = df[['intercept','yield','uniform_ordered','accepted_first_rank','had_enrolled_sib','ell_status','homeless_status','es_school','ms_school','orientation_rsvp','virtual_event_attended','in_person_event_attended',\n",
    "            'scholar_grade','commute_time',\n",
    "            'log_commute','log_commute_square','log_commute_third',\n",
    "            'school','utm_source_bucketing']]\n",
    "    df = df.dropna()\n",
    "    return df\n",
    "\n",
    "def set_predictors(current_predictors):\n",
    "    predictors = current_predictors\n",
    "    return predictors\n",
    "\n",
    "def set_target(current_target):\n",
    "    target = current_target\n",
    "    return target\n",
    "\n",
    "def setup_training_data(df, predictors, target):\n",
    "    X = df[predictors]\n",
    "    Y = df[target]\n",
    "    return X, Y\n",
    "\n",
    "def test_train_split(X, Y):\n",
    "    X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.25, random_state=0)\n",
    "    return X_train, X_test, y_train, y_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>intercept</th>\n",
       "      <th>yield</th>\n",
       "      <th>uniform_ordered</th>\n",
       "      <th>accepted_first_rank</th>\n",
       "      <th>had_enrolled_sib</th>\n",
       "      <th>ell_status</th>\n",
       "      <th>homeless_status</th>\n",
       "      <th>es_school</th>\n",
       "      <th>ms_school</th>\n",
       "      <th>orientation_rsvp</th>\n",
       "      <th>virtual_event_attended</th>\n",
       "      <th>in_person_event_attended</th>\n",
       "      <th>scholar_grade</th>\n",
       "      <th>commute_time</th>\n",
       "      <th>log_commute</th>\n",
       "      <th>log_commute_square</th>\n",
       "      <th>log_commute_third</th>\n",
       "      <th>school</th>\n",
       "      <th>utm_source_bucketing</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>46.616667</td>\n",
       "      <td>3.841958</td>\n",
       "      <td>14.760642</td>\n",
       "      <td>56.709770</td>\n",
       "      <td>SA Bronx 4</td>\n",
       "      <td>Organic / No Tracking</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>46.616667</td>\n",
       "      <td>3.841958</td>\n",
       "      <td>14.760642</td>\n",
       "      <td>56.709770</td>\n",
       "      <td>SA Bronx 4</td>\n",
       "      <td>Organic / No Tracking</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>80.066667</td>\n",
       "      <td>4.382860</td>\n",
       "      <td>19.209458</td>\n",
       "      <td>84.192360</td>\n",
       "      <td>SA Washington Heights</td>\n",
       "      <td>Google Branded Search</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>5.100000</td>\n",
       "      <td>1.629241</td>\n",
       "      <td>2.654425</td>\n",
       "      <td>4.324696</td>\n",
       "      <td>SA Far Rockaway</td>\n",
       "      <td>Organic / No Tracking</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>34.000000</td>\n",
       "      <td>3.526361</td>\n",
       "      <td>12.435219</td>\n",
       "      <td>43.851064</td>\n",
       "      <td>SA Bronx 2</td>\n",
       "      <td>Organic / No Tracking</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19269</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>8.800000</td>\n",
       "      <td>2.174752</td>\n",
       "      <td>4.729545</td>\n",
       "      <td>10.285586</td>\n",
       "      <td>SA Washington Heights</td>\n",
       "      <td>Organic / No Tracking</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19270</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>27.166667</td>\n",
       "      <td>3.301991</td>\n",
       "      <td>10.903143</td>\n",
       "      <td>36.002076</td>\n",
       "      <td>SA Bergen Beach</td>\n",
       "      <td>Organic / No Tracking</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19271</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>12.283333</td>\n",
       "      <td>2.508243</td>\n",
       "      <td>6.291285</td>\n",
       "      <td>15.780073</td>\n",
       "      <td>SA Bronx 1</td>\n",
       "      <td>Google Branded Search</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19272</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>28.200000</td>\n",
       "      <td>3.339322</td>\n",
       "      <td>11.151071</td>\n",
       "      <td>37.237017</td>\n",
       "      <td>SA Harlem 3</td>\n",
       "      <td>Organic / No Tracking</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19273</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>14.016667</td>\n",
       "      <td>2.640247</td>\n",
       "      <td>6.970905</td>\n",
       "      <td>18.404911</td>\n",
       "      <td>SA Harlem 3</td>\n",
       "      <td>Organic / No Tracking</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>19273 rows × 19 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "       intercept  yield  uniform_ordered  accepted_first_rank  \\\n",
       "0              1      0                0                    1   \n",
       "1              1      0                0                    1   \n",
       "2              1      0                0                    0   \n",
       "3              1      1                1                    1   \n",
       "4              1      0                0                    0   \n",
       "...          ...    ...              ...                  ...   \n",
       "19269          1      0                0                    1   \n",
       "19270          1      0                0                    1   \n",
       "19271          1      0                0                    1   \n",
       "19272          1      0                1                    1   \n",
       "19273          1      0                0                    1   \n",
       "\n",
       "       had_enrolled_sib  ell_status  homeless_status  es_school  ms_school  \\\n",
       "0                     0           1                0          1          0   \n",
       "1                     0           1                0          1          0   \n",
       "2                     0           1                0          1          0   \n",
       "3                     0           1                0          1          0   \n",
       "4                     0           0                0          1          0   \n",
       "...                 ...         ...              ...        ...        ...   \n",
       "19269                 0           0                0          1          0   \n",
       "19270                 0           1                0          1          0   \n",
       "19271                 0           1                0          1          0   \n",
       "19272                 0           0                0          1          0   \n",
       "19273                 0           1                0          1          0   \n",
       "\n",
       "       orientation_rsvp  virtual_event_attended  in_person_event_attended  \\\n",
       "0                     0                       0                         0   \n",
       "1                     0                       0                         0   \n",
       "2                     0                       0                         0   \n",
       "3                     1                       0                         0   \n",
       "4                     0                       0                         0   \n",
       "...                 ...                     ...                       ...   \n",
       "19269                 1                       0                         0   \n",
       "19270                 0                       0                         0   \n",
       "19271                 0                       0                         0   \n",
       "19272                 0                       0                         0   \n",
       "19273                 0                       0                         0   \n",
       "\n",
       "       scholar_grade  commute_time  log_commute  log_commute_square  \\\n",
       "0                  1     46.616667     3.841958           14.760642   \n",
       "1                  1     46.616667     3.841958           14.760642   \n",
       "2                  3     80.066667     4.382860           19.209458   \n",
       "3                  3      5.100000     1.629241            2.654425   \n",
       "4                  2     34.000000     3.526361           12.435219   \n",
       "...              ...           ...          ...                 ...   \n",
       "19269              2      8.800000     2.174752            4.729545   \n",
       "19270              2     27.166667     3.301991           10.903143   \n",
       "19271              0     12.283333     2.508243            6.291285   \n",
       "19272              0     28.200000     3.339322           11.151071   \n",
       "19273              0     14.016667     2.640247            6.970905   \n",
       "\n",
       "       log_commute_third                 school   utm_source_bucketing  \n",
       "0              56.709770             SA Bronx 4  Organic / No Tracking  \n",
       "1              56.709770             SA Bronx 4  Organic / No Tracking  \n",
       "2              84.192360  SA Washington Heights  Google Branded Search  \n",
       "3               4.324696        SA Far Rockaway  Organic / No Tracking  \n",
       "4              43.851064             SA Bronx 2  Organic / No Tracking  \n",
       "...                  ...                    ...                    ...  \n",
       "19269          10.285586  SA Washington Heights  Organic / No Tracking  \n",
       "19270          36.002076        SA Bergen Beach  Organic / No Tracking  \n",
       "19271          15.780073             SA Bronx 1  Google Branded Search  \n",
       "19272          37.237017            SA Harlem 3  Organic / No Tracking  \n",
       "19273          18.404911            SA Harlem 3  Organic / No Tracking  \n",
       "\n",
       "[19273 rows x 19 columns]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = _get_table_data('train_basetable.csv')\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "target = set_target(['yield'])\n",
    "predictors = set_predictors(['uniform_ordered','accepted_first_rank','had_enrolled_sib',\n",
    "                            'orientation_rsvp','virtual_event_attended','in_person_event_attended'])\n",
    "X, y =  setup_training_data(df, predictors, target)\n",
    "X = torch.tensor(X.values, dtype=torch.float32)\n",
    "y = torch.tensor(y.values, dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[iteration 0050] loss: 71834256.0000\n",
      "[iteration 0100] loss: 71754288.0000\n",
      "[iteration 0150] loss: 71752720.0000\n",
      "[iteration 0200] loss: 71752696.0000\n",
      "[iteration 0250] loss: 71752696.0000\n",
      "[iteration 0300] loss: 71752680.0000\n",
      "[iteration 0350] loss: 71752680.0000\n",
      "[iteration 0400] loss: 71752680.0000\n",
      "[iteration 0450] loss: 71752680.0000\n",
      "[iteration 0500] loss: 71752680.0000\n",
      "[iteration 0550] loss: 71752680.0000\n",
      "[iteration 0600] loss: 71752680.0000\n",
      "[iteration 0650] loss: 71752680.0000\n",
      "[iteration 0700] loss: 71752680.0000\n",
      "[iteration 0750] loss: 71752680.0000\n",
      "[iteration 0800] loss: 71752680.0000\n",
      "[iteration 0850] loss: 71752680.0000\n",
      "[iteration 0900] loss: 71752680.0000\n",
      "[iteration 0950] loss: 71752680.0000\n",
      "[iteration 1000] loss: 71752680.0000\n",
      "[iteration 1050] loss: 71752680.0000\n",
      "[iteration 1100] loss: 71752680.0000\n",
      "[iteration 1150] loss: 71752680.0000\n",
      "[iteration 1200] loss: 71752680.0000\n",
      "[iteration 1250] loss: 71752680.0000\n",
      "[iteration 1300] loss: 71752680.0000\n",
      "[iteration 1350] loss: 71752680.0000\n",
      "[iteration 1400] loss: 71752680.0000\n",
      "[iteration 1450] loss: 71752688.0000\n",
      "[iteration 1500] loss: 71752680.0000\n",
      "Learned parameters:\n",
      "weight [[ 9.63660263e-09 -1.39678882e-08  1.74582127e-10  7.05604508e-09\n",
      "   1.45985410e-08  1.03843085e-08]]\n",
      "bias [0.2616095]\n"
     ]
    }
   ],
   "source": [
    "# Regression model\n",
    "linear_reg_model = PyroModule[nn.Linear](6, 1)\n",
    "\n",
    "# Define loss and optimize\n",
    "loss_fn = torch.nn.MSELoss(reduction='sum')\n",
    "optim = torch.optim.Adam(linear_reg_model.parameters(), lr=0.05)\n",
    "num_iterations = 1500\n",
    "\n",
    "\n",
    "def train(x_data,y_data):\n",
    "    # run the model forward on the data\n",
    "    y_pred = linear_reg_model(x_data).squeeze(-1)\n",
    "    # calculate the mse loss\n",
    "    loss = loss_fn(y_pred, y_data)\n",
    "    # initialize gradients to zero\n",
    "    optim.zero_grad()\n",
    "    # backpropagate\n",
    "    loss.backward()\n",
    "    # take a gradient step\n",
    "    optim.step()\n",
    "    return loss\n",
    "\n",
    "for j in range(num_iterations):\n",
    "    loss = train(X,y)\n",
    "    if (j + 1) % 50 == 0:\n",
    "        print(\"[iteration %04d] loss: %.4f\" % (j + 1, loss.item()))\n",
    "\n",
    "# Inspect learned parameters\n",
    "print(\"Learned parameters:\")\n",
    "for name, param in linear_reg_model.named_parameters():\n",
    "    print(name, param.data.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit = df.copy()\n",
    "fit[\"mean\"] = linear_reg_model(X).detach().cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyro.nn import PyroSample\n",
    "\n",
    "\n",
    "class BayesianRegression(PyroModule):\n",
    "    def __init__(self, in_features, out_features):\n",
    "        super().__init__()\n",
    "        self.linear = PyroModule[nn.Linear](in_features, out_features)\n",
    "        self.linear.weight = PyroSample(dist.Normal(0., 1.).expand([out_features, in_features]).to_event(2))\n",
    "        self.linear.bias = PyroSample(dist.Normal(0., 10.).expand([out_features]).to_event(1))\n",
    "\n",
    "    def forward(self, x, y=None):\n",
    "        sigma = pyro.sample(\"sigma\", dist.Uniform(0., 10.))\n",
    "        mean = self.linear(x).squeeze(-1)\n",
    "        with pyro.plate(\"data\", x.shape[0]):\n",
    "            obs = pyro.sample(\"obs\", dist.Normal(mean, sigma), obs=y)\n",
    "        return mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyro.infer.autoguide import AutoDiagonalNormal\n",
    "\n",
    "model = BayesianRegression(6, 1)\n",
    "guide = AutoDiagonalNormal(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyro.infer import SVI, Trace_ELBO\n",
    "\n",
    "\n",
    "adam = pyro.optim.Adam({\"lr\": 0.03})\n",
    "svi = SVI(model, guide, adam, loss=Trace_ELBO())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "at site \"obs\", invalid log_prob shape\n  Expected [19273], actual [19273, 19273]\n  Try one of the following fixes:\n  - enclose the batched tensor in a with pyro.plate(...): context\n  - .to_event(...) the distribution being sampled\n  - .permute() data dimensions",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[24], line 4\u001b[0m\n\u001b[1;32m      1\u001b[0m pyro\u001b[38;5;241m.\u001b[39mclear_param_store()\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m j \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(num_iterations):\n\u001b[1;32m      3\u001b[0m     \u001b[38;5;66;03m# calculate the loss and take a gradient step\u001b[39;00m\n\u001b[0;32m----> 4\u001b[0m     loss \u001b[38;5;241m=\u001b[39m \u001b[43msvi\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstep\u001b[49m\u001b[43m(\u001b[49m\u001b[43mX\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      5\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m j \u001b[38;5;241m%\u001b[39m \u001b[38;5;241m100\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m      6\u001b[0m         \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m[iteration \u001b[39m\u001b[38;5;132;01m%04d\u001b[39;00m\u001b[38;5;124m] loss: \u001b[39m\u001b[38;5;132;01m%.4f\u001b[39;00m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m%\u001b[39m (j \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m1\u001b[39m, loss \u001b[38;5;241m/\u001b[39m \u001b[38;5;28mlen\u001b[39m(df)))\n",
      "File \u001b[0;32m~/Documents/yield_forecast/edna-data-science-talend-models/yield_dev/.venv/lib/python3.12/site-packages/pyro/infer/svi.py:145\u001b[0m, in \u001b[0;36mSVI.step\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    143\u001b[0m \u001b[38;5;66;03m# get loss and compute gradients\u001b[39;00m\n\u001b[1;32m    144\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m poutine\u001b[38;5;241m.\u001b[39mtrace(param_only\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m) \u001b[38;5;28;01mas\u001b[39;00m param_capture:\n\u001b[0;32m--> 145\u001b[0m     loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mloss_and_grads\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mguide\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    147\u001b[0m params \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mset\u001b[39m(\n\u001b[1;32m    148\u001b[0m     site[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mvalue\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39munconstrained() \u001b[38;5;28;01mfor\u001b[39;00m site \u001b[38;5;129;01min\u001b[39;00m param_capture\u001b[38;5;241m.\u001b[39mtrace\u001b[38;5;241m.\u001b[39mnodes\u001b[38;5;241m.\u001b[39mvalues()\n\u001b[1;32m    149\u001b[0m )\n\u001b[1;32m    151\u001b[0m \u001b[38;5;66;03m# actually perform gradient steps\u001b[39;00m\n\u001b[1;32m    152\u001b[0m \u001b[38;5;66;03m# torch.optim objects gets instantiated for any params that haven't been seen yet\u001b[39;00m\n",
      "File \u001b[0;32m~/Documents/yield_forecast/edna-data-science-talend-models/yield_dev/.venv/lib/python3.12/site-packages/pyro/infer/trace_elbo.py:140\u001b[0m, in \u001b[0;36mTrace_ELBO.loss_and_grads\u001b[0;34m(self, model, guide, *args, **kwargs)\u001b[0m\n\u001b[1;32m    138\u001b[0m loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0.0\u001b[39m\n\u001b[1;32m    139\u001b[0m \u001b[38;5;66;03m# grab a trace from the generator\u001b[39;00m\n\u001b[0;32m--> 140\u001b[0m \u001b[43m\u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mmodel_trace\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mguide_trace\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_get_traces\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mguide\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m    141\u001b[0m \u001b[43m    \u001b[49m\u001b[43mloss_particle\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msurrogate_loss_particle\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_differentiable_loss_particle\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    142\u001b[0m \u001b[43m        \u001b[49m\u001b[43mmodel_trace\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mguide_trace\u001b[49m\n\u001b[1;32m    143\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    144\u001b[0m \u001b[43m    \u001b[49m\u001b[43mloss\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mloss_particle\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnum_particles\u001b[49m\n",
      "File \u001b[0;32m~/Documents/yield_forecast/edna-data-science-talend-models/yield_dev/.venv/lib/python3.12/site-packages/pyro/infer/elbo.py:237\u001b[0m, in \u001b[0;36mELBO._get_traces\u001b[0;34m(self, model, guide, args, kwargs)\u001b[0m\n\u001b[1;32m    235\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    236\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_particles):\n\u001b[0;32m--> 237\u001b[0m         \u001b[38;5;28;01myield\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_get_trace\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mguide\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/Documents/yield_forecast/edna-data-science-talend-models/yield_dev/.venv/lib/python3.12/site-packages/pyro/infer/trace_elbo.py:57\u001b[0m, in \u001b[0;36mTrace_ELBO._get_trace\u001b[0;34m(self, model, guide, args, kwargs)\u001b[0m\n\u001b[1;32m     52\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_get_trace\u001b[39m(\u001b[38;5;28mself\u001b[39m, model, guide, args, kwargs):\n\u001b[1;32m     53\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     54\u001b[0m \u001b[38;5;124;03m    Returns a single trace from the guide, and the model that is run\u001b[39;00m\n\u001b[1;32m     55\u001b[0m \u001b[38;5;124;03m    against it.\u001b[39;00m\n\u001b[1;32m     56\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 57\u001b[0m     model_trace, guide_trace \u001b[38;5;241m=\u001b[39m \u001b[43mget_importance_trace\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     58\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mflat\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmax_plate_nesting\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mguide\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\n\u001b[1;32m     59\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     60\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m is_validation_enabled():\n\u001b[1;32m     61\u001b[0m         check_if_enumerated(guide_trace)\n",
      "File \u001b[0;32m~/Documents/yield_forecast/edna-data-science-talend-models/yield_dev/.venv/lib/python3.12/site-packages/pyro/infer/enum.py:80\u001b[0m, in \u001b[0;36mget_importance_trace\u001b[0;34m(graph_type, max_plate_nesting, model, guide, args, kwargs, detach)\u001b[0m\n\u001b[1;32m     78\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m site \u001b[38;5;129;01min\u001b[39;00m model_trace\u001b[38;5;241m.\u001b[39mnodes\u001b[38;5;241m.\u001b[39mvalues():\n\u001b[1;32m     79\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m site[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtype\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msample\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m---> 80\u001b[0m         \u001b[43mcheck_site_shape\u001b[49m\u001b[43m(\u001b[49m\u001b[43msite\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmax_plate_nesting\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     81\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m site \u001b[38;5;129;01min\u001b[39;00m guide_trace\u001b[38;5;241m.\u001b[39mnodes\u001b[38;5;241m.\u001b[39mvalues():\n\u001b[1;32m     82\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m site[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtype\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msample\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n",
      "File \u001b[0;32m~/Documents/yield_forecast/edna-data-science-talend-models/yield_dev/.venv/lib/python3.12/site-packages/pyro/util.py:437\u001b[0m, in \u001b[0;36mcheck_site_shape\u001b[0;34m(site, max_plate_nesting)\u001b[0m\n\u001b[1;32m    433\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m actual_size, expected_size \u001b[38;5;129;01min\u001b[39;00m zip_longest(\n\u001b[1;32m    434\u001b[0m     \u001b[38;5;28mreversed\u001b[39m(actual_shape), \u001b[38;5;28mreversed\u001b[39m(expected_shape), fillvalue\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m\n\u001b[1;32m    435\u001b[0m ):\n\u001b[1;32m    436\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m expected_size \u001b[38;5;241m!=\u001b[39m \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m \u001b[38;5;129;01mand\u001b[39;00m expected_size \u001b[38;5;241m!=\u001b[39m actual_size:\n\u001b[0;32m--> 437\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    438\u001b[0m             \u001b[38;5;124m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m  \u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(\n\u001b[1;32m    439\u001b[0m                 [\n\u001b[1;32m    440\u001b[0m                     \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mat site \u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m, invalid log_prob shape\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(site[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mname\u001b[39m\u001b[38;5;124m\"\u001b[39m]),\n\u001b[1;32m    441\u001b[0m                     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mExpected \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m, actual \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(expected_shape, actual_shape),\n\u001b[1;32m    442\u001b[0m                     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTry one of the following fixes:\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m    443\u001b[0m                     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m- enclose the batched tensor in a with pyro.plate(...): context\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m    444\u001b[0m                     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m- .to_event(...) the distribution being sampled\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m    445\u001b[0m                     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m- .permute() data dimensions\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m    446\u001b[0m                 ]\n\u001b[1;32m    447\u001b[0m             )\n\u001b[1;32m    448\u001b[0m         )\n\u001b[1;32m    450\u001b[0m \u001b[38;5;66;03m# Check parallel dimensions on the left of max_plate_nesting.\u001b[39;00m\n\u001b[1;32m    451\u001b[0m enum_dim \u001b[38;5;241m=\u001b[39m site[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124minfer\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m_enumerate_dim\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[0;31mValueError\u001b[0m: at site \"obs\", invalid log_prob shape\n  Expected [19273], actual [19273, 19273]\n  Try one of the following fixes:\n  - enclose the batched tensor in a with pyro.plate(...): context\n  - .to_event(...) the distribution being sampled\n  - .permute() data dimensions"
     ]
    }
   ],
   "source": [
    "pyro.clear_param_store()\n",
    "for j in range(num_iterations):\n",
    "    # calculate the loss and take a gradient step\n",
    "    loss = svi.step(X, y)\n",
    "    if j % 100 == 0:\n",
    "        print(\"[iteration %04d] loss: %.4f\" % (j + 1, loss / len(df)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}