<jupyter_start><jupyter_text>Example of grouped regressionsIn this section, we want to demanstrate a slightly advanced example for using Pandas grouped transformation for performing many ordinary least square model fits in parallel. We reuse the weather data and try to predict the temperature of all stations with a very simple model per station.<jupyter_code>%matplotlib inline<jupyter_output><empty_output><jupyter_text>1 Load DataFirst we load data of a single year.<jupyter_code>storageLocation = "s3://dimajix-training/data/weather"
from pyspark.sql.functions import *
from pyspark.sql.types import *


rawWeatherData = spark.read.text(storageLocation + "/2003")
weather_all = rawWeatherData.select(
    substring(col("value"), 5, 6).alias("usaf"),
    substring(col("value"), 11, 5).alias("wban"),
    to_timestamp(substring(col("value"), 16, 12), "yyyyMMddHHmm").alias("timestamp"),
    to_timestamp(substring(col("value"), 16, 12), "yyyyMMddHHmm")
    .cast("long")
    .alias("ts"),
    substring(col("value"), 42, 5).alias("report_type"),
    substring(col("value"), 61, 3).alias("wind_direction"),
    substring(col("value"), 64, 1).alias("wind_direction_qual"),
    substring(col("value"), 65, 1).alias("wind_observation"),
    (substring(col("value"), 66, 4).cast("float") / lit(10.0)).alias("wind_speed"),
    substring(col("value"), 70, 1).alias("wind_speed_qual"),
    (substring(col("value"), 88, 5).cast("float") / lit(10.0)).alias("air_temperature"),
    substring(col("value"), 93, 1).alias("air_temperature_qual"),
)<jupyter_output><empty_output><jupyter_text>2 Analysis of one stationFirst we only analyse a single station, just to check our approach and the expressiveness of our model. It won't be a very good fit, but it will be good enough for our needs to demonstrate the concept.So first we pick a single station, and we also only keep those records with a valid temeprature measurement.<jupyter_code>weather_single = weather_all.where("usaf='954920' and wban='99999'").cache()
pdf = weather_single.where(weather_single.air_temperature_qual == 1).toPandas()
pdf.limit(10)<jupyter_output><empty_output><jupyter_text>2.1 Create Feature SpaceOur model will simply predict the temperature depending on the time and day of year. We use sin and cos of with a day-wide period and a year-wide period as features for fitting the model.<jupyter_code>import math

import numpy as np


seconds_per_day = 24 * 60 * 60
seconds_per_year = 365 * seconds_per_day

# Add sin and cos as features for fitting
pdf['daily_sin'] = np.sin(pdf['ts'] / seconds_per_day * 2.0 * math.pi)
pdf['daily_cos'] = np.cos(pdf['ts'] / seconds_per_day * 2.0 * math.pi)
pdf['yearly_sin'] = np.sin(pdf['ts'] / seconds_per_year * 2.0 * math.pi)
pdf['yearly_cos'] = np.cos(pdf['ts'] / seconds_per_year * 2.0 * math.pi)

# Make a plot, just to check how it looks like
pdf[0:200].plot(
    x='timestamp', y=['daily_sin', 'daily_cos', 'air_temperature'], figsize=[16, 6]
)<jupyter_output><empty_output><jupyter_text>2.2 Fit modelNow that we have the temperature and some features, we fit a simple model.<jupyter_code>import statsmodels.api as sm


# define target variable y
y = pdf['air_temperature']
# define feature variables X
X = pdf[['ts', 'daily_sin', 'daily_cos', 'yearly_sin', 'yearly_cos']]
X = sm.add_constant(X)
# fit model
model = sm.OLS(y, X).fit()

# perform prediction
pdf['pred'] = model.predict(X)

# Make a plot of real temperature vs predicted temperature
pdf[0:200].plot(x='timestamp', y=['pred', 'air_temperature'], figsize=[16, 6])<jupyter_output><empty_output><jupyter_text>2.3 Inspect ModelNow let us inspect the model, in order to find a way to store it in a Pandas DataFrame<jupyter_code>model.params
type(model.params)<jupyter_output><empty_output><jupyter_text>Finally let us create a Pandas DataFrame from the model parameters. This code snippet will be needed later when we want to parallelize the fitting for different weather stations using Spark.<jupyter_code>x_columns = X.columns
pd.DataFrame([[model.params[i] for i in x_columns]], columns=x_columns)<jupyter_output><empty_output><jupyter_text>3 Perform OLS for all stationsNow we want to create a model for all stations. First we filter the data again, such that we only have valid temperature measurements.<jupyter_code>valid_weather = weather_all.filter(weather_all.air_temperature_qual == 1)<jupyter_output><empty_output><jupyter_text>3.1 Feature extractionNow we generate the same features, but this time we use Spark instead of Pandas operations. This simplifies later model fitting.<jupyter_code>import math


seconds_per_day = 24 * 60 * 60
seconds_per_year = 365 * seconds_per_day

features = valid_weather.select(
    valid_weather.usaf,
    valid_weather.wban,
    valid_weather.air_temperature,
    valid_weather.ts,
    lit(1.0).alias('const'),
    sin(valid_weather.ts * 2.0 * math.pi / seconds_per_day).alias('daily_sin'),
    cos(valid_weather.ts * 2.0 * math.pi / seconds_per_day).alias('daily_cos'),
    sin(valid_weather.ts * 2.0 * math.pi / seconds_per_year).alias('yearly_sin'),
    cos(valid_weather.ts * 2.0 * math.pi / seconds_per_year).alias('yearly_cos'),
)

features.limit(10).toPandas()<jupyter_output><empty_output><jupyter_text>3.2 Fit ModelsNow we use a Spark Pandas grouped UDF in order to fit models for all weather stations in parallel. This is done by defining a Pandas Gruped Map UDF, which receives all values of a single weather station as its input and returns a Pandas data frame containing one row with all model parameters. This UDF is then used in a grouped map aggregation using `groupby(...).apply(...)`.<jupyter_code>group_columns = ['usaf', 'wban']
y_column = 'air_temperature'
x_columns = ['ts', 'const', 'daily_sin', 'daily_cos', 'yearly_sin', 'yearly_cos']
schema = features.select(*group_columns, *x_columns).schema


@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def ols(pdf):
    # Extract grouping information from appropriate columns
    group = [pdf[g].iloc[0] for g in group_columns]
    # Extract target variable
    y = pdf[y_column]
    # Extract predictor variables
    X = pdf[x_columns]
    # Create model using Python statsmodel package to fit y to input variables x
    model = sm.OLS(y, X).fit()
    # Create a Pandas data frame with one row containing the grouping columns and all model parameters
    return pd.DataFrame(
        [group + [model.params[i] for i in x_columns]], columns=group_columns + x_columns
    )


# Now fit model for all weather stations in parallel using Spark
models = features.groupby(weather_all.usaf, weather_all.wban).apply(ols).cache()
models.limit(10).toPandas()<jupyter_output><empty_output><jupyter_text>3.3 Inspect and compare resultsNow let's pick the same station again, and compare the model to the original model.<jupyter_code>models.where("usaf='954920' and wban='99999'").toPandas()
model.params<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11804750869061413, 'fraction_numerical': 0.02085747392815759, 'mean_word_length': 5.501883239171375, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 45, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '230800', 'license': 'MIT', 'n_tokens_mistral': 2337, 'n_tokens_neox': 2160, 'n_words': 662, 'path': 'spark-training/spark-python/jupyter-advanced-udf/Grouped Regression - Full.ipynb', 'repo_name': 'Code360In/spark-code-examples', 'size': '6904'}