Save a Python function as a model

FoundryML provides a native wrapper for incorporating Python functions operating on a single pandas DataFrame. This enables arbitrary processing and rapid prototyping, as well as implementation of non-row-wise models that must operate on an entire dataset at once, such as some forecasting models, simulations, and optimizations.

Python functions meeting the criteria below can be annotated with the foundry_ml.function_stages.pandas_function_stage decorator. This generates a stage, which can be combined with other stages (if needed) to form a Python model. Models containing these stages are fully compatible with Batch and Live deployments, and can be invoked on pandas or pyspark DataFrames (though in the latter case, the processing will still occur locally in pandas).

The high level requirements for a function to be decorated are below:

1. The decorated function's signature should match.

   (data: pandas.DataFrame, params: NamedTuple) -> pandas.DataFrame
   

2. data is a required kwarg (keyword argument) and must be a pandas DataFrame.
3. params is a required kwarg and must be a NamedTuple.

Example: Implementing an SIR Model

This tutorial demonstrates how to implement an SIR (Susceptible, Infected, Recovered) model in Foundry. This is a simple version of a larger class of "compartmentalized" models that are used to forecast the spread of a pandemic. For this example, we will focus on re-implementing the simple SIR model ↗ provided in the SciPy documentation in the FoundryML ecosystem using the pandas_function_stage decorator.

The following code implements the core functionality of the model:

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import numpy as np
from scipy.integrate import odeint

# Total population, N.
N = 1000
# Initial number of infected and recovered individuals, I0 and R0.
I0, R0 = 1, 0
# Everyone else, S0, is susceptible to infection initially.
S0 = N - I0 - R0
# Contact rate, beta, and mean recovery rate, gamma, (in 1/days).
beta, gamma = 0.2, 1./10
# A grid of time points (in days)
t = np.linspace(0, 160, 160)

# The SIR model differential equations.
def deriv(y, t, N, beta, gamma):
    S, I, R = y
    dSdt = -beta * S * I / N
    dIdt = beta * S * I / N - gamma * I
    dRdt = gamma * I
    return dSdt, dIdt, dRdt

# Initial conditions vector
y0 = S0, I0, R0
# Integrate the SIR equations over the time grid, t.
ret = odeint(deriv, y0, t, args=(N, beta, gamma))
S, I, R = ret.T

Getting Started

Create a new Code Workbook and skip the introductory dialog to import datasets into the environment. Then, create a new Transform and select Python as the language. We also assume that there is a dataset of states and corresponding populations, with the fields id: string, population: int. Note, if you don't have a dataset, you can also define a Pandas DataFrame which you can use for testing with the following code:

Copied!
1
2
3
4
5
6
7
import pandas as pd

input_df: pd.DataFrame = pd.DataFrame([
                                ['TX', 29000000],
                                ['CO', 5800000],
                                ['NY', 8400000]],
                                columns=['id', 'population'])

Use the pandas_function_stage decorator

In order to integrate this model with FoundryML, we will to use the pandas_function_stage decorator. This decorator will prepare our function such that it can be wrapped in a Foundry ML Stage, which we can then use to create a Foundry ML Model.

To do this, we'll first define a NamedTuple to apply some typing to our input parameters.

Copied!
1
2
3
4
5
6
7
8
from typing import NamedTuple

class SIRParameters(NamedTuple):
    n_days: int = 2
    I0: int = 1
    R0: int = 0
    beta: float = 0.2
    gamma: float = 0.1

We'll then create a forecast method to run the forecast for a single state and use sir_model to apply this function to each row of the states DataFrame.

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
from scipy.integrate import odeint
import pandas as pd

def sir_model(data: pd.DataFrame, params: SIRParameters) -> pd.DataFrame:
    rows = data.to_dict('rows')
    results = []
    for state in rows:
        id = state['id']
        population = state['population']
        state_df = forecast(population, params)
        state_df['state'] = id
        results.append(state_df)
    return pd.concat(results, ignore_index=True)

def forecast(n_population, parameters):
    n_days = parameters.n_days
    I0 = parameters.I0
    R0 = parameters.R0
    S0 = n_population - I0 - R0

    beta = parameters.beta
    gamma = parameters.gamma
    t = np.linspace(0, n_days, n_days)

    y0 = S0, I0, R0
    ret = odeint(deriv, y0, t, args=(n_days, beta, gamma))
    S, I, R = ret.T
    return pd.DataFrame({
        'susceptible': S,
        'infected': I,
        'recovered': R,
        'n_days': range(n_days),
    })

def deriv(y, t, N, beta, gamma):
    S, I, R = y
    dSdt = -beta * S * I / N
    dIdt = beta * S * I / N - gamma * I
    dRdt = gamma * I
    return dSdt, dIdt, dRdt

Lastly, decorate your sir_model function with the pandas_function_stage decorator.

Copied!
1
2
3
4
5
from foundry_ml.function_stages import pandas_function_stage

@pandas_function_stage()
def sir_model(data: pd.DataFrame, params: SIRParameters) -> pd.DataFrame:
    ...

With the function property decorated, you can now save the model as a Foundry ML Model.

Copied!
1
2
3
4
from foundry_ml import Model, Stage

def model():
    return Model(Stage(sir_model))

You can execute the model like any other Model with the transform function.

Copied!
1
2
3
def execute_model(model, states_df: pd.DataFrame):
    result_no_overrides = model.transform(states_df)
    result_with_override = model.transform(states_df, params={'n_days': 10})

Execute your model in a live deployment

If you are serving your model through a live deployment, you can execute the deployment using the live deployment inference API:

curl --http2 -H "Content-Type: application/json" -H "Authorization: <$BEARER_TOKEN>" -d '{"requestData":[{"id":"TX","population":29000000}, {"id":"CO","population":5800000}], "requestParams":{"params":{"num_days":10}}}' --request POST $STACK_URL/foundry-ml-live/api/inference/transform/$DEPLOYMENT_RID