# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Day 4: Intro to Grama
#
# - Worked example
# - Grama motivation
# - Grama elements
# - Exploratory model analysis
#

# +
import grama as gr

import numpy as np
import scipy as sp
import pandas as pd
import matplotlib.pyplot as plt
from plotnine import *
DF = gr.Intention()

# Set figure options
plt.rcParams['figure.figsize'] = [6, 6] # Need square aspect ratio for gradients to appear normal
plt.rcParams['figure.dpi'] = 100 # 200 e.g. is really fine, but slower

# -

# ## Worked Example: Fitting a model

# Load data
#

# +
## Load built-in dataset
from grama.data import df_trajectory_windowed

(
    df_trajectory_windowed
    >> ggplot(aes("x", "y"))
    + geom_point()
)
# -

# Suppose we want to make predictions on where the projectile will land; we'll need a *model*:

# +
## Load built-in Grama model
from grama.models import make_trajectory_linear
md_traj = make_trajectory_linear()

## Print info about this model
md_traj.printpretty()

# -

# This is a model for a projectile's trajectory using Newton's laws and a linear drag model. The model takes in inputs `u0, v0, tau` (parameters) and the time `t` and outputs `x, y` coordinates.
#
# Making a prediction with uninformed guesses will be terrible:

# +
## Need to set parameters for prediction:
u0 = 20 # initial horizontal velocity
v0 = 20 # initial vertical velocity  
tau = 5 # time constant (drag)

## Make a prediction with the model
df_prediction = gr.eval_df(
    md_traj,
    df=gr.df_make(u0=u0, v0=v0, tau=tau, t=np.linspace(0, 2))
)

## Visualize
(
    df_trajectory_windowed
    >> ggplot(aes("x", "y"))
    + geom_point()
    + geom_line(data=df_prediction, color="blue")
)
# -

# Use *nonlinear least squares* (NLS) to **optimize** the parameter values:
#

# +
## Fit the model
md_fit = gr.fit_nls(
    df_trajectory_windowed,
    md=md_traj,
    method="SLSQP",
    seed=101,
    verbose=False,
)
md_fit.printpretty()

## Make prediction with fitted model
df_pred_fit = gr.eval_df(
    md_fit,
    df=gr.df_make(t=np.linspace(0, 4.5))
)

## Visualize
(
    df_trajectory_windowed
    >> ggplot(aes("x", "y"))
    + geom_point()
    + geom_line(data=df_pred_fit, color="salmon")
)

# -

# ## Grama Motivation
#
# ---
#
# Idea: Have one computational object that you can use to encapsulate *many* of the assumptions and choices in quantitative modeling.
#
# Some example modeling choices:
# - Select the model inputs
# - Specify bounds for the inputs
# - Simplify the physics
# - Fit a distribution for inputs
#
# Grama is a software package to help make model building and analysis more transparent, convenient, and understandable.
#
# Modeling in two phases:
# - Model building
# - Model analysis
#

# ### Model Building
#
# ---
#
# Make a blank model
#

## Create a blank model
md = gr.Model("Base model")
## Print the details
md.printpretty()

# Make a model with a function

## Create a blank model
md = gr.Model("Base model")
## Add a function
md = gr.comp_function(
    md, 
    fun=lambda X: X[0], # f(x) = x
    var=["x"],          # Inputs: x
    out=["f"],          # Outputs: f
)
## Print the details
md.printpretty()

# Make a model with function and variable bounds

## Create a blank model
md = gr.Model("Base model")
## Add a function
md = gr.comp_function(
    md, 
    fun=lambda X: X[0],
    var=["x"],
    out=["f"],
)
## Add a bound
md = gr.comp_bounds(
    md,
    x=(-1, +1),
)
## Print the details
md.printpretty()

# ### Example: RLC circuit
#
# **Model Functions**
#
# Parallel RLC circuit: here are some key performance indicators of the circuit
#
# $$\omega_0 = \sqrt{\frac{1}{LC}}$$
#
# $$Q = \omega_0 RC.$$
#
# **Model Domain**
#
# Decide on a range of values for $R, L, C$ to test.

# +
## Implement RLC Grama model
md_RLC = (
    gr.Model("RLC Circuit")
    ## Add the natural frequency
    >> gr.cp_vec_function(
        fun=lambda df: gr.df_make(
            omega0=np.sqrt(1 / df.L / df.C)
        ),
        name="natural frequency",
        var=["L", "C"],
        out=["omega0"],
    )
    ## Add the Q factor
    >> gr.cp_vec_function(
        fun=lambda df: gr.df_make(
            Q=df.omega0 * df.R * df.C
        ),
        name="quality factor",
        var=["omega0", "R", "C"],
        out=["Q"]
    )
    ## Set bounds for the input variables
    >> gr.cp_bounds(
        R=(1e-3, 1e0),  # resistance [ohm]
        L=(1e-9, 1e-3), # inductance [H]
        C=(1e-3, 10),   # capacitance [F]
    )
)

md_RLC.printpretty()
# -

# **Remember**: You can always use `md.printpretty()` to inspect a model; see its inputs, outputs, function names, bounds/distribution.
#

# ## Model Analysis
#
# Evaluate the model at specified `R,L,C` values:
#

gr.eval_df(
    md_RLC,
    df=gr.df_make(R=1000, L=0.1, C=0.1)
)


# Fit the model to find `R,L,C` values for specified `omega0, Q`.

df_rlc = gr.eval_nls(
    md_RLC,
    df_data=gr.df_make(omega0=10, Q=1),
    n_restart=10
)
df_rlc


# Check that those values give the desired values
#

gr.eval_df(md_RLC, df=df_rlc)


# ## Grama Elements
#
# ---
#
# Grama considers *data* and *models*. Data are organized into DataFrames, which are handled by the Pandas package.
#

df_example = pd.DataFrame(dict(
    x=[1, 2, 3],
    y=[0.1, 0.2, 0.3],
    z=["a", "b", "c"]
))
df_example


# DataFrames are useful because they're more human-readable than arrays. Each column has a name, so we can access specific columns with `df.variable`:

df_example.x


# There are four fundamental Grama [verbs](https://py-grama.readthedocs.io/en/latest/source/language.html#verbs):
#
# | Verb Type | Prefix (Short) | In | Out |
# |---|---|---|---|
# | Compose | `comp_` (`cp_`) | `md` | `md` |
# | Evaluate | `eval_` (`ev_`) | `md` | `df` |
# | Fit | `fit_` (`ft_`) | `df` | `md` |
# | Transform | `tran_` (`tf_`) | `df` | `df` |
# | Plot | `plot_` (`pt_`) | `df` | (Plot) |
#

# ### Compose
#
# Used primarily to build up a model
#

## Create a blank model
md = gr.Model("Base model")
## Add a function
md = gr.comp_function(
    ## Take in function; will return modified
    md, 
    fun=lambda X: X[0],
    var=["x"],
    out=["f"],
)
## Print the details
md.printpretty()


# ### Evaluate
#
# Used to generate data from a model
#

df_result = gr.eval_df(
    ## Model to evaluate
    md,
    ## DataFrame at which to evaluate
    df=gr.df_make(x=[0, 1, 2])
)
df_result


# ### Fit
#
# Used to derive a model from data
#
# First, set up a scenario with data and a model to fit:
#

# +
## 
from grama.models import make_trajectory_linear
from grama.data import df_trajectory_windowed

md_trajectory = make_trajectory_linear()
md_trajectory.printpretty()
# -

# Fit the model
#

md_fit = gr.fit_nls(df_trajectory_windowed, md_trajectory)
md_fit.printpretty()

# ### Transform
#
# Used to transform data
#

df_trajectory_windowed.head()


# Estimate time derivatives with finite differences
#

(
    df_trajectory_windowed
    >> gr.tf_mutate(
        # Estimate horizontal velocity
        #    (x1 - x0) / (t1 - t0)
        dxdt=(DF.x - gr.lag(DF.x)) / (DF.t - gr.lag(DF.t)),
        dydt=(DF.y - gr.lag(DF.y)) / (DF.t - gr.lag(DF.t)),
    )
    >> gr.tf_head()
)

# ## Exploratory Model Analysis
#
# ---
#
# Grama is useful for *exploratory model analysis*; making sense of how a model behaves with respect to its inputs. Let's look at a simple model to build intuition.
#
#

# $$f(x, a) = a \exp(x)$$

# +
md_exponential = (
    ## Start an empty model
    gr.Model("Exponential model")
    ## Add in the function
    >> gr.cp_vec_function(
        fun=lambda df: gr.df_make(
            f=df.a * np.exp(df.x)
        ),
        var=["a", "x"],
        out=["f"],
        name="Exponential"
    )
    ## Add some bounds
    >> gr.cp_bounds(
        a=(-1, +1),
        x=(-1, +1),
    )
)

md_exponential.printpretty()
# -

# Let's investigate the model with a *sinew* plot
#

(
    md_exponential
    >> gr.ev_sinews(df_det="swp", seed=101)
    >> gr.pt_auto()
)

# - input `a` has a linear effect on output `f`
# - input `x` has an exponential effect on output `f`
#   - direction is affected by `a`
#

# Sinew plots especially useful for exploring a new model
#

df_results = (
    md_RLC
    >> gr.ev_sinews(df_det="swp", n_density=20, n_sweeps=5, seed=101)
    >> gr.pt_auto()
)


# - input `R` has a positive, linear effect on `Q`
# - input `R` has zero effect on `omega0`
# - input `C` has a positive, diminishing effect on `Q`
# - input `C` has a negative, diminishing effect on `omega0`
# - input `L` has a positive, diminishing effect on `Q`
# - input `L` has a negative, diminishing effect on `omega0`
#

StarcoderdataJupyterScriptsDedupFiltered