## Introduction to FLARE: Fast Learning of Atomistic Rare Events
Jonathan Vandermause (jonathan_vandermause@g.harvard.edu)

**Learning objectives:**
  * Train 2+3-body Gaussian process models on _ab initio_ force data.
  * Use the uncertainties of the GP to train a force field on the fly.

Computing the properties of real materials with quantum mechanical accuracy is extremely expensive. One of the most efficient explicitly quantum mechanical tools we have for the job, density functional theory (or DFT), has a computational cost that scales cubically with the number of atoms in the system, and is therefore limited in practice to at most a few hundred atoms (with some hope of scaling further with orbital free DFT methods).

One of the most popular strategies for getting around this is to simply ignore the hard part of the problem---the electrons---and to instead try to model the system with an _empirical interatomic potential_, which expresses the potential energy as a sum over local, atom-centered contributions that depend  only on atomic coordinates. The resulting model scales only linearly with the number of atoms, making it possible to simulate the behavior of hundreds of thousands or even millions of atoms over microsecond timescales. The problem then becomes: how do you design a good potential? Is it possible to find a fast empirical potential with the accuracy of DFT?

FLARE is an open-source software that uses Bayesian machine learning to try to bridge the gap between accurate but slow quantum methods (like DFT) and fast but inaccurate classical methods (like empirical interatomic potentials). The idea is to train fast potentials on accurate DFT data, and ideally to do so in an automatic, closed loop fashion, so that a wide class of materials and material compositions can be efficiently explored. The key tool FLARE uses to accomplish this is Gaussian process (GP) regression, an elegant framework for defining probability distributions over functions. In this tutorial, we'll explore the use of GPs to model interactions between atoms in materials.

If you're interested in learning more about the FLARE code, check out our [paper](https://www.nature.com/articles/s41524-020-0283-z), [GitHub page](https://github.com/mir-group/flare), and [code documentation](https://flare.readthedocs.io/en/latest/).

### Installation

We can install FLARE and its dependencies here in Google Colab. (This will take about a minute.)

```python
! pip install --upgrade mir-flare
```

Let's check that it worked by trying to import FLARE:

```python
import flare
```

If you don't see an error, you're all set for the tutorial! Let's also go ahead and import all the modules we'll need now:

```python
from flare import gp, struc, output, predict, md, otf, env,\
  otf_parser
from flare.kernels import mc_simple
from flare.utils import md_helper

import numpy as np
import time
import matplotlib.pyplot as plt
import pickle
from ase.visualize import view
```

### Machine learned force fields

#### Training data

Let's start by downloading some _ab initio_ force data to play with.

```python
# download ab initio MD data
! wget https://zenodo.org/record/3688843/files/AgI_data.zip?download=1

# unzip the folder
! unzip AgI_data.zip?download=1
```

The folder "AgI_data" contains data from an _ab initio_ molecular dynamics simulation of the fast-ion conductor silver iodide (AgI), where _ab initio_ means "from the beginning" or "from first principles". In other words, we simulate how the atoms move by calculating the quantum mechanical forces on the ions at every timestep using DFT. To give you a sense of how expensive that is, the ~2500 timesteps in this simulation (about 12 picoseconds of simulation time) required 2 days of wall time on 256 cpus. (And there were only 32 atoms in the simulation!)

Let's load the positions, forces, cell, and species of the atoms in the simulation.

```python
# load AIMD training data
data_directory = 'AgI_data/'
species = np.load(data_directory + 'species.npy')  # atomic numbers of the atoms
positions = np.load(data_directory + 'positions.npy')  # in angstrom (A)
cell = np.load(data_directory + 'cell.npy')  # 3x3 array of cell vectors (in A)
forces = np.load(data_directory + 'forces.npy')  # in eV/A
```

#### Training a GP model.

Let's train a GP model on the data. We'll first initialize a Gaussian process object, which involves choosing a kernel function and its gradient, an initial set of hyperparameters, and the cutoff radii of local environments. We'll use a 2+3-body kernel, which compares local environments by comparing the pairs and triplets of atoms inside the environments.

```python
# create a 2+3-body gaussian process object
kernels = ['twobody', 'threebody']
component = 'mc'
hyps = np.array([0.1, 0.1, 0.1, 2.0, 0.5])  # initial (bad) choice of hyps
cutoffs = {'twobody': 7.0, 'threebody': 5.5}  # cutoff radii in A
maxiter = 100  # max number of hyperparameter optimziation steps

gp_model = gp.GaussianProcess(
  kernels=kernels,
  component=component,
  hyps=hyps,
  cutoffs=cutoffs,
  maxiter=50
)
```

Let's put a couple of structures into the training set.

```python
# put a few snapshots in the training set
snapshots = [500, 1500]
for snapshot in snapshots:
    # create flare structure
    training_positions = positions[snapshot]
    training_forces = forces[snapshot]
    training_structure = struc.Structure(cell, species, training_positions)

    # add the structure to the training set of the GP
    gp_model.update_db(training_structure, training_forces)

gp_model.set_L_alpha()
```

The "set_L_alpha" method updates the covariance matrix of the GP, computes the alpha vector used to make predictions, and computes the log marginal likelihood of the training data. Let's see what that looks like:

```python
print(gp_model.likelihood)
```

    -491.75669961275526

This is too low, suggesting that our initial guess for the hyperparameters was not a good one.

#### Task: Find hyperparameters that give a positive log marginal likelihood.

Hint: For the 2+3-body kernel, the hyperparameters are in the following order:
* Signal variance of the 2-body kernel (in eV)
* Length scale of the 2-body kernel (in A)
* Signal variance of the 3-body kernel (in eV)
* Length scale of the 3-body kernel (in A)
* Noise hyperparameter (in eV/A)

What is a plausible length scale for this problem? Note that machine learned force fields typically have errors in the range 50-200 meV/A, and the energy of a pair or triplet is much less than the total local energy assigned to an atom.

```python
# your code here

# hint: reset the hyperparameters with gp_model.hyps = *new hyps*, then use
# set_L_alpha to recompute the likelihood
```

#### Solution

A reasonable guess for the length scale is 1 A, and since errors are  often around 0.1 eV/A, we'll choose that as our noise level. The signal variances require some tuning, but we should expect them to be significantly less than 1 eV, with the triplet contribution significantly smaller than the pair contribution.

```python
gp_model.hyps = np.array([0.01, 1, 0.001, 1, 0.2])
gp_model.set_L_alpha()
print(gp_model.likelihood)
```

    9.500574018895804

#### Optimizing the hyperparameters rigorously

The above was an exercise in manual hyperparameter tuning, which is tedious and should be avoided if possible. Because we can compute the gradient of the likelihood with respect to the hyperparameters, we can instead use a more principled gradient descent approach to find the best set of hyperparameters.

```python
# optimize the hyperparameters (this will take some time!)
gp_model.train(print_progress=True)
```

In case you don't want to wait, here are some good values:

```python
gp_model.hyps = np.array([1.59612454e-02, 5.70104540e-01, 6.01290125e-04,
                          9.17243358e-01, 1.09666317e-01])
gp_model.set_L_alpha()
print(gp_model.likelihood)
```

    71.16991245736048

#### Calculating the learning curve

Let's get a feel for how much data the model needs by computing the learning curve, i.e. the performance on a validation set as a function of the number of training points. Let's create training and validation structures drawn from the AIMD simulation.

```python
# choose a training snapshot
training_snapshot = 500
training_positions = positions[training_snapshot]
training_forces = forces[training_snapshot]
training_structure = struc.Structure(cell, species, training_positions)

# choose a validation snapshot
validation_snapshot = 2300
validation_positions = positions[validation_snapshot]
validation_forces = forces[validation_snapshot]
validation_structure = struc.Structure(cell, species, validation_positions)
```

We now loop over the atomic environments in the training structure and add them one by one to the training set of the GP model. After adding an environment, we predict all the forces on the validation structure and compute the MAE. We'll also time the prediction step to get a feel for how the cost of GP predictions depends on the size of the training set.

```python
# reset the gp with hyperparameters fixed to the optimized values
hyps_final = gp_model.hyps
gp_model = gp.GaussianProcess(
  kernels=kernels,
  component=component,
  hyps=hyps,
  cutoffs=cutoffs,
  maxiter=50
)
gp_model.hyps = hyps_final

# add atomic environments one by one to the training set
n_atoms = 32  # number of atoms in the structure
validation_errors = np.zeros(n_atoms)
prediction_times = np.zeros(n_atoms)
validation_count = 0

print('computing validation errors...')
for n, atom in enumerate(range(n_atoms)):
    # add the current atomic environment to the training set
    gp_model.update_db(training_structure, training_forces,
                        custom_range=[atom])
    gp_model.set_L_alpha()

    # predict the forces on the validation structure
    time0 = time.time()
    pred_forces, stds = \
        predict.predict_on_structure(validation_structure, gp_model)
    time1 = time.time()

    mae = np.mean(np.abs(pred_forces - validation_forces))
    validation_errors[validation_count] = mae
    prediction_times[validation_count] = time1 - time0
    validation_count += 1

    print(mae)
```

    computing validation errors...
    0.315073761272364
    0.29588515439108337
    0.3234313322939279
    0.29323785059806706
    0.20129532813447323
    0.19131142668143372
    0.1827709765647784
    0.17610478136490326
    0.16902120241436183
    0.17466172091548085
    0.16829757387913633
    0.16823160886145053
    0.1797266370933175
    0.17570045468126091
    0.1778018841113225
    0.16181457527423726
    0.15221658747132247
    0.151463002400943
    0.15330262545396417
    0.14366990594229426
    0.13955981466966869
    0.1386958938185947
    0.13337032134291107
    0.13856470639127907
    0.13934854695295343
    0.1384439054077334
    0.1376182149699538
    0.1398980197468981
    0.13878505117526346
    0.13947592439212986
    0.14140188721074384
    0.14095157040791634

Let's make a plot of the results.

```python
plt.plot(validation_errors)
plt.xlabel('number of training environments')
plt.ylabel('mae on validation structure (eV/$\AA$)')
plt.show()

plt.plot(prediction_times)
plt.xlabel('number of training environments')
plt.ylabel('prediction wall time (s)')
plt.show()
```

Notice that the prediction time grows roughly linearly with the number of training environments.

### Learning a force field on the fly

#### The Lennard Jones Potential

In production FLARE runs, we make calls to a DFT solver like Quantum Espresso, VASP, or CP2K whenever the uncertainty on a force component is unacceptably high. In principle, the forces can come from anywhere, and FLARE allows users to write a custom solver that returns forces on an arbitrary input structure.

Here, we'll define a simple custom solver that returns Lennard Jones forces, and then try to reconstruct the potential on the fly.

The Lennard Jones potential takes the following form:
\begin{equation}
V_{\text{LJ}} = 4\epsilon \left[ \left(\frac{\sigma}{r} \right)^{12} - \left(\frac{\sigma}{r}\right)^{6} \right].
\end{equation}

Let's define our custom solver. To do this, we need to create a class with two methods:

* parse_dft_input: Takes the name of an input file, and returns the positions, species, cell, and masses of the atoms specified in the input file.

* run_dft_par: Takes a FLARE structure and an optional dictionary of keywords, and returns the forces on the atoms.

We'll make a multi-species Lennard Jones function, which allows different values of $\sigma$ and $\epsilon$ for different pairs of species. We'll store these as 2-D numpy arrays, so that, e.g., the $\sigma$ value assigned to species 1 and 3 is stored in element (1, 3) of this array.

The derivative of the Lennard Jones potential with respect to $r$ is:

\begin{equation}
\frac{dV_{\text{LJ}}}{dr} = 4\epsilon \left[ -\frac{12 \sigma^{12}}{r^{13}} +\frac{6 \sigma^6}{r^7} \right].
\end{equation}

```python
def get_LJ_forces(structure, lj_parameters):
  """Calculate multicomponent Lennard Jones forces on a structure of atoms.
  dft_kwargs is assumed to be a dictionary containing the cutoff, an ordered
  list of species, and arrays containing epsilon and sigma values."""

  cutoff = lj_parameters['cutoff']
  epsilons = lj_parameters['epsilons']
  sigmas = lj_parameters['sigmas']
  spec_list = lj_parameters['species']

  forces = np.zeros((structure.nat, 3))
  # Loop over atoms in the structure.
  for m in range(structure.nat):
      # Create atomic environment.
      environment = env.AtomicEnvironment(structure, m, np.array([cutoff]))
      ind1 = spec_list.index(environment.ctype)

      # Loop over atoms in the environment to compute the total force.
      for n in range(len(environment.etypes)):
          ind2 = spec_list.index(environment.etypes[n])
          eps = epsilons[ind1, ind2]
          sig = sigmas[ind1, ind2]

          # Compute LJ force.
          bond_vals = environment.bond_array_2[n]
          r = bond_vals[0]

          dE_dr = 4 * eps * (-12 * sig ** 12 / (r ** 13) +
                             6 * sig ** 6 / (r ** 7))

          forces[m, 0] += dE_dr * bond_vals[1]
          forces[m, 1] += dE_dr * bond_vals[2]
          forces[m, 2] += dE_dr * bond_vals[3]

  return forces
```

We can use our Lennard Jones force function to define a custom module.

```python
# create lj module
class lj_module:
  def parse_dft_input(file_name):
      """We assume the input is a pickled dictionary containing positions
      (in angstrom), species (as a list of strings), cell (as a 3x3 matrix of
      cell vectors), and masses (as a dictionary assigning each species a mass
      in AMU)."""

      input_file = open(file_name, 'rb')
      struc_dict = pickle.load(input_file)
      input_file.close()

      # Convert masses to MD units (energy = eV, length = A, )
      masses = struc_dict['masses']
      mass_convert = 0.000103642695727
      for species in masses:
          masses[species] *= mass_convert

      return struc_dict['positions'], struc_dict['species'], \
          struc_dict['cell'], masses

  def run_dft_par(dft_input=None, structure=None, dft_loc=None, n_cpus=None,
                npool=None, mpi=None, dft_kwargs=None, dft_out=None):
    return get_LJ_forces(structure, dft_kwargs)
```

```python
# create dictionary of lj parameters
cutoff = 5.
epsilons = np.array([[3, 2.5], [2.5, 3.]])
sigmas = np.array([[2.0, 1.9], [1.9, 2.1]])
spec_list = [47, 53]  # silver and iodine atomic numbers

lj_params = {'cutoff': cutoff, 'epsilons': epsilons, 'sigmas': sigmas,
             'species': spec_list}
```

```python
# create structures of 2 atoms, and plot the force on the second atom
cell = np.eye(3) * 1000
seps = np.arange(2, 5, 0.01)
specs = [[47, 47], [47, 53], [53, 53]]
store_frcs = np.zeros((3, len(seps)))
for m, sep in enumerate(seps):
    pos = np.array([[0, 0, 0], [sep, 0, 0]])
    for n, spec_curr in enumerate(specs):
        struc_curr = struc.Structure(cell, spec_curr, pos)
        frcs = \
            lj_module.run_dft_par(structure=struc_curr, dft_kwargs=lj_params)
        store_frcs[n, m] = frcs[1, 0]

plt.plot(seps, store_frcs[0], label='sig = 2.0, eps = 3.0')
plt.plot(seps, store_frcs[1], label='sig = 1.5, eps = 1.8')
plt.plot(seps, store_frcs[2], label='sig = 1.0, eps = 1.0')
plt.xlim(2, 4)
plt.ylim(-4, 4)
plt.xlabel('separation ($\AA$)')
plt.ylabel('force (eV/$\AA$)')
plt.legend()
plt.show()
```

#### Perform an on-the-fly training simulation.

#### Step 1: Set up the initial structure.

 Let's create a bcc (body centered cubic) Lennard Jones crystal.

```python
# create bcc unit cell
alat = 2.54
unit_cell = np.eye(3) * alat

# define bcc positions
unit_positions = np.array([[0, 0, 0],
                           [1/2, 1/2, 1/2]]) * alat

# make a supercell
sc_size = 3
positions = md_helper.get_supercell_positions(sc_size, unit_cell, unit_positions)
cell = unit_cell * sc_size

# jitter positions to give nonzero force on first frame
for atom_pos in positions:
    for coord in range(3):
        atom_pos[coord] += (2*np.random.random()-1) * 0.05

# create initial structure
species = ['Ag', 'I'] * sc_size ** 3
struc_curr = struc.Structure(cell, species, positions)

# create pseudo input file
mass_dictionary = {'Ag': 108, 'I': 127}
input_dictionary = {'positions': positions, 'species': species, 'cell': cell,
                    'masses': mass_dictionary}
input_file_name = 'lj.in'
with open(input_file_name, 'wb') as input_file:
    pickle.dump(input_dictionary, input_file)
```

```python
# view the structure
view(struc_curr.to_ase_atoms(), viewer='x3d')
```

<html>

 <head>

  <title>ASE atomic visualization</title>

  <link rel="stylesheet" type="text/css"

   href="https://www.x3dom.org/x3dom/release/x3dom.css">

  </link>

 </head>

 <body>

  <X3D>

   <Scene>

    <Transform translation="0.01 -0.05 -0.01">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="1.26 1.30 1.27">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="0.04 -0.00 2.56">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="1.29 1.29 3.83">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="0.03 -0.00 5.08">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="1.26 1.23 6.31">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="0.04 2.52 0.04">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="1.23 3.85 1.30">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="-0.01 2.56 2.57">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="1.28 3.77 3.85">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="0.01 2.54 5.09">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="1.26 3.82 6.40">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="-0.00 5.11 0.03">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="1.28 6.34 1.27">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="0.01 5.11 2.55">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="1.30 6.34 3.83">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="-0.02 5.11 5.06">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="1.30 6.30 6.37">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="2.54 0.03 0.04">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="3.79 1.27 1.31">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="2.57 0.03 2.56">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="3.76 1.26 3.80">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="2.49 -0.00 5.13">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="3.78 1.31 6.31">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="2.56 2.54 0.00">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="3.79 3.82 1.29">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="2.50 2.51 2.55">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="3.85 3.77 3.76">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="2.57 2.57 5.06">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="3.79 3.80 6.35">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="2.54 5.04 -0.02">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="3.80 6.39 1.24">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="2.58 5.08 2.54">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="3.77 6.40 3.82">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="2.55 5.06 5.07">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="3.85 6.30 6.32">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="5.10 -0.05 -0.00">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="6.32 1.26 1.23">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="5.06 -0.02 2.56">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="6.38 1.26 3.82">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="5.10 -0.02 5.10">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="6.37 1.24 6.30">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="5.10 2.55 0.01">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="6.35 3.76 1.30">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="5.10 2.51 2.49">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="6.31 3.83 3.81">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="5.09 2.58 5.06">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="6.38 3.82 6.35">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="5.07 5.05 -0.01">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="6.36 6.37 1.23">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="5.12 5.11 2.52">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="6.39 6.37 3.79">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="5.12 5.09 5.12">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.753 0.753 0.753" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.45">

      </Sphere>

     </Shape>

    </Transform>

    <Transform translation="6.31 6.37 6.36">

     <Shape>

      <Appearance>

       <Material diffuseColor="0.580 0.000 0.580" specularColor="0.5 0.5 0.5">

       </Material>

      </Appearance>

      <Sphere radius="1.39">

      </Sphere>

     </Shape>

    </Transform>

   </Scene>

  </X3D>

 </body>

</html>

#### Step 2: Set up a GP model.

```python
# create gaussian process model
kernels = ['twobody']
component = 'mc'
hyps = np.array([0.1, 1., 0.06])
hyp_labels = ['Sigma', 'Length Scale', 'Noise']
cutoffs = {'twobody': 5.0}
maxiter = 50

gp_model = gp.GaussianProcess(
  kernels=kernels,
  component=component,
  hyps=hyps,
  hyp_labels=hyp_labels,
  cutoffs=cutoffs,
  maxiter=50
)
```

#### Step 3: Set up an OTF training object.

```python
# create otf object
dt = 0.001  # time step (ps)
number_of_steps = 50
dft_loc = None  # path to dft executable would usually go here
std_tolerance_factor = -0.01  # 10 meV/A
init_atoms = [0, 25, 50]  # initial atoms added to the training set
max_atoms_added = 5  # number of atoms added when dft is called
freeze_hyps = 5  # no hyperparameter optimization after this many updates

# rescale the temperature halfway through the simulation
rescale_steps = [10]  # rescale at step 10
rescale_temps = [1000]

otf_model = otf.OTF(
    # MD arguments
    dt,
    number_of_steps,
    rescale_steps=rescale_steps,
    rescale_temps=rescale_temps,
    # FLARE arguments
    gp=gp_model,
    # OTF arguments
    std_tolerance_factor=std_tolerance_factor,
    init_atoms=init_atoms,
    max_atoms_added=max_atoms_added,
    freeze_hyps=freeze_hyps,
    # DFT arguments
    force_source=lj_module,
    dft_loc=dft_loc,
    dft_input=input_file_name,
    dft_kwargs=lj_params
    )
```

#### Step 4: Perform the simulation!

```python
# perform flare run (this will take about 2.5 minutes)
otf_model.run()
```

#### Parsing the OTF output file.

```python
# parse output file
output_file = 'otf_run.out'
otf_trajectory = otf_parser.OtfAnalysis(output_file)

# plot temperature vs. simulation time
times = otf_trajectory.times
temps = otf_trajectory.thermostat['temperature']
dft_times = otf_trajectory.dft_times[1:]  # exclude t = 0

for n, dft_time in enumerate(dft_times):
    otf_ind = times.index(dft_time)
    plt.plot(dft_times[n], temps[otf_ind], 'kx')

plt.plot(times, temps)
plt.xlabel('time (ps)')
plt.ylabel('temperature (K)')
plt.show()
```

#### Comparing the learned model against ground truth.

```python
# create structures of 2 atoms, and plot the force on the second atom
cell = np.eye(3) * 1000
seps = np.arange(2, 5, 0.01)
specs = [[47, 47], [47, 53], [53, 53]]
store_frcs = np.zeros((3, len(seps)))
gp_frcs = np.zeros((3, len(seps)))
gp_stds = np.zeros((3, len(seps)))
for m, sep in enumerate(seps):
    pos = np.array([[0, 0, 0], [sep, 0, 0]])
    for n, spec_curr in enumerate(specs):
        struc_curr = struc.Structure(cell, spec_curr, pos)
        env_curr = env.AtomicEnvironment(struc_curr, 1, np.array([cutoff]))
        frcs = \
            lj_module.run_dft_par(structure=struc_curr, dft_kwargs=lj_params)
        store_frcs[n, m] = frcs[1, 0]

        # predict the x component of the force
        pred, var = gp_model.predict(env_curr, 1)

        gp_frcs[n, m] = pred
        gp_stds[n, m] = np.sqrt(var)

# plot GP predictions vs ground truth
cols = ['b', 'r', 'g']
for n in range(3):
    plt.plot(seps, store_frcs[n], color=cols[n], linestyle='-')
    plt.plot(seps, gp_frcs[n], color=cols[n], linestyle='--')
    plt.fill_between(seps, gp_frcs[n] + 3 * gp_stds[n],
                     gp_frcs[n] - 3 * gp_stds[n],
                     color=cols[n], alpha=0.4)

plt.xlabel('separation ($\AA$)')
plt.ylabel('force (eV/$\AA$)')
plt.xlim(2, 3.4)
plt.ylim(-10, 15)
plt.show()
```

```python

```


# Orthogonal Random Forest and Causal Forest: Use Cases and Examples

Causal Forests and Generalized Random Forests are a flexible method for estimating treatment effect heterogeneity with Random Forests. Orthogonal Random Forest (ORF) combines orthogonalization, a technique that effectively removes the confounding effect in two-stage estimation, with generalized random forests. Due to the orthogonalization aspect of this method, the ORF performs especially well in the presence of high-dimensional confounders. For more details, see [this paper](https://arxiv.org/abs/1806.03467) or the [EconML docummentation](https://econml.azurewebsites.net/).

The EconML SDK implements the following OrthoForest variants:

* DMLOrthoForest: suitable for continuous or discrete treatments

* DROrthoForest: suitable for discrete treatments

* CausalForest: suitable for both discrete and continuous treatments

In this notebook, we show the performance of the ORF on synthetic and observational data.

## Notebook Contents

1. [Example Usage with Continuous Treatment Synthetic Data](#1.-Example-Usage-with-Continuous-Treatment-Synthetic-Data)
2. [Example Usage with Binary Treatment Synthetic Data](#2.-Example-Usage-with-Binary-Treatment-Synthetic-Data)
3. [Example Usage with Multiple Treatment Synthetic Data](#3.-Example-Usage-with-Multiple-Treatment-Synthetic-Data)
4. [Example Usage with Real Continuous Treatment Observational Data](#4.-Example-Usage-with-Real-Continuous-Treatment-Observational-Data)

```python
import econml
```

```python
# Main imports
from econml.orf import DMLOrthoForest, DROrthoForest
from econml.dml import CausalForestDML
from econml.sklearn_extensions.linear_model import WeightedLassoCVWrapper, WeightedLasso, WeightedLassoCV

# Helper imports
import numpy as np
from itertools import product
from sklearn.linear_model import Lasso, LassoCV, LogisticRegression, LogisticRegressionCV
import matplotlib.pyplot as plt

%matplotlib inline
```

# 1. Example Usage with Continuous Treatment Synthetic Data

## 1.1 DGP
We use the data generating process (DGP) from [here](https://arxiv.org/abs/1806.03467). The DGP is described by the following equations:

\begin{align}
T =& \langle W, \beta\rangle + \eta, & \;\eta \sim \text{Uniform}(-1, 1)\\
Y =& T\cdot \theta(X) + \langle W, \gamma\rangle + \epsilon, &\; \epsilon \sim \text{Uniform}(-1, 1)\\
W \sim& \text{Normal}(0,\, I_{n_w})\\
X \sim& \text{Uniform}(0,1)^{n_x}
\end{align}

where $W$ is a matrix of high-dimensional confounders and $\beta, \gamma$ have high sparsity.

For this DGP,
\begin{align}
\theta(x) = \exp(2\cdot x_1).
\end{align}

```python
# Treatment effect function
def exp_te(x):
    return np.exp(2*x[0])
```

```python
# DGP constants
np.random.seed(123)
n = 1000
n_w = 30
support_size = 5
n_x = 1
# Outcome support
support_Y = np.random.choice(range(n_w), size=support_size, replace=False)
coefs_Y = np.random.uniform(0, 1, size=support_size)
epsilon_sample = lambda n: np.random.uniform(-1, 1, size=n)
# Treatment support
support_T = support_Y
coefs_T = np.random.uniform(0, 1, size=support_size)
eta_sample = lambda n: np.random.uniform(-1, 1, size=n)

# Generate controls, covariates, treatments and outcomes
W = np.random.normal(0, 1, size=(n, n_w))
X = np.random.uniform(0, 1, size=(n, n_x))
# Heterogeneous treatment effects
TE = np.array([exp_te(x_i) for x_i in X])
T = np.dot(W[:, support_T], coefs_T) + eta_sample(n)
Y = TE * T + np.dot(W[:, support_Y], coefs_Y) + epsilon_sample(n)

# ORF parameters and test data
subsample_ratio = 0.3
lambda_reg = np.sqrt(np.log(n_w) / (10 * subsample_ratio * n))
X_test = np.array(list(product(np.arange(0, 1, 0.01), repeat=n_x)))
```

## 1.2. Train Estimator

**Note:** The models in the final stage of the estimation (``model_T_final``, ``model_Y_final``) need to support sample weighting.

If the models of choice do not support sample weights (e.g. ``sklearn.linear_model.LassoCV``), the ``econml`` packages provides a convenient wrapper for these models ``WeightedModelWrapper`` in order to allow sample weights.

```python
est = DMLOrthoForest(
    n_trees=1000, min_leaf_size=5,
    max_depth=50, subsample_ratio=subsample_ratio,
    model_T=Lasso(alpha=lambda_reg),
    model_Y=Lasso(alpha=lambda_reg),
    model_T_final=WeightedLasso(alpha=lambda_reg),
    model_Y_final=WeightedLasso(alpha=lambda_reg),
    global_residualization=False,
    random_state=123)
```

To use the built-in confidence intervals constructed via Bootstrap of Little Bags, we can specify `inference="blb"` at `fit` time or leave the default `inference='auto'` which will automatically use the Bootstrap of Little Bags.

```python
est.fit(Y, T, X=X, W=W, inference="blb")
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:   21.6s
    [Parallel(n_jobs=-1)]: Done 176 tasks      | elapsed:   22.6s
    [Parallel(n_jobs=-1)]: Done 816 tasks      | elapsed:   25.6s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:   26.5s finished
    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    0.0s
    [Parallel(n_jobs=-1)]: Done 368 tasks      | elapsed:    1.7s
    [Parallel(n_jobs=-1)]: Done 984 tasks      | elapsed:    4.7s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:    4.7s finished

    <econml.orf._ortho_forest.DMLOrthoForest at 0x1a7d2b58a58>

```python
# Calculate treatment effects
treatment_effects = est.effect(X_test)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:   21.6s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:   23.9s finished

```python
# Calculate default (95%) confidence intervals for the test data
te_lower, te_upper = est.effect_interval(X_test)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  18 tasks      | elapsed:    3.3s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:    7.6s finished

```python
res = est.effect_inference(X_test)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    3.4s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:    7.4s finished

```python
res.summary_frame().head()
```

      point_estimate
      stderr
      zstat
      pvalue
      ci_lower
      ci_upper

      X

      0
      1.161
      0.183
      6.339
      0.0
      0.802
      1.520

      1
      1.171
      0.177
      6.628
      0.0
      0.825
      1.518

      2
      1.182
      0.171
      6.925
      0.0
      0.847
      1.516

      3
      1.192
      0.165
      7.228
      0.0
      0.869
      1.515

      4
      1.202
      0.160
      7.533
      0.0
      0.890
      1.515

```python
res.population_summary()
```

Uncertainty of Mean Point Estimate

  mean_point	 stderr_mean	 zstat	 pvalue	 ci_mean_lower	 ci_mean_upper

     3.179	      0.287	    11.06	   0.0	      2.616	         3.742

Distribution of Point Estimate

  std_point	 pct_point_lower	 pct_point_upper

    1.715	        1.187	           6.276

Total Variance of Point Estimate

  stderr_point	 ci_point_lower	 ci_point_upper

      1.739	         1.079	          6.525

Note: The stderr_mean is a conservative upper bound.

Similarly we can estimate effects and get confidence intervals and inference results using a `CausalForest`.

```python
est2 = CausalForestDML(model_t=Lasso(alpha=lambda_reg),
                       model_y=Lasso(alpha=lambda_reg),
                       n_estimators=4000, min_samples_leaf=5,
                       max_depth=50,
                       verbose=0, random_state=123)
est2.tune(Y, T, X=X, W=W)
est2.fit(Y, T, X=X, W=W)
treatment_effects2 = est2.effect(X_test)
te_lower2, te_upper2 = est2.effect_interval(X_test, alpha=0.01)
```

## 1.3. Performance Visualization

```python
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.title("ContinuousOrthoForest")
plt.plot(X_test, treatment_effects, label='ORF estimate')
expected_te = np.array([exp_te(x_i) for x_i in X_test])
plt.plot(X_test[:, 0], expected_te, 'b--', label='True effect')
plt.fill_between(X_test[:, 0], te_lower, te_upper, label="95% BLB CI", alpha=0.3)
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.subplot(1, 2, 2)
plt.title("CausalForest")
plt.plot(X_test, treatment_effects2, label='ORF estimate')
expected_te = np.array([exp_te(x_i) for x_i in X_test])
plt.plot(X_test[:, 0], expected_te, 'b--', label='True effect')
plt.fill_between(X_test[:, 0], te_lower2, te_upper2, label="95% BLB CI", alpha=0.3)
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.show()
```

# 2. Example Usage with Binary Treatment Synthetic Data

## 2.1. DGP
We use the following DGP:

\begin{align}
T \sim & \text{Bernoulli}\left(f(W)\right), &\; f(W)=\sigma(\langle W, \beta\rangle + \eta), \;\eta \sim \text{Uniform}(-1, 1)\\
Y = & T\cdot \theta(X) + \langle W, \gamma\rangle + \epsilon, & \; \epsilon \sim \text{Uniform}(-1, 1)\\
W \sim & \text{Normal}(0,\, I_{n_w}) & \\
X \sim & \text{Uniform}(0,\, 1)^{n_x}
\end{align}

where $W$ is a matrix of high-dimensional confounders, $\beta, \gamma$ have high sparsity and $\sigma$ is the sigmoid function.

For this DGP,
\begin{align}
\theta(x) = \exp( 2\cdot x_1 ).
\end{align}

```python
# DGP constants
np.random.seed(1234)
n = 1000
n_w = 30
support_size = 5
n_x = 1
# Outcome support
support_Y = np.random.choice(range(n_w), size=support_size, replace=False)
coefs_Y = np.random.uniform(0, 1, size=support_size)
epsilon_sample = lambda n: np.random.uniform(-1, 1, size=n)
# Treatment support
support_T = support_Y
coefs_T = np.random.uniform(0, 1, size=support_size)
eta_sample = lambda n: np.random.uniform(-1, 1, size=n)

# Generate controls, covariates, treatments and outcomes
W = np.random.normal(0, 1, size=(n, n_w))
X = np.random.uniform(0, 1, size=(n, n_x))
# Heterogeneous treatment effects
TE = np.array([exp_te(x_i) for x_i in X])
# Define treatment
log_odds = np.dot(W[:, support_T], coefs_T) + eta_sample(n)
T_sigmoid = 1/(1 + np.exp(-log_odds))
T = np.array([np.random.binomial(1, p) for p in T_sigmoid])
# Define the outcome
Y = TE * T + np.dot(W[:, support_Y], coefs_Y) + epsilon_sample(n)

# ORF parameters and test data
subsample_ratio = 0.4
X_test = np.array(list(product(np.arange(0, 1, 0.01), repeat=n_x)))
```

## 2.2. Train Estimator

```python
est = DROrthoForest(
    n_trees=200, min_leaf_size=10,
    max_depth=30, subsample_ratio=subsample_ratio,
    propensity_model = LogisticRegression(C=1/(X.shape[0]*lambda_reg), penalty='l1', solver='saga'),
    model_Y = Lasso(alpha=lambda_reg),
    propensity_model_final=LogisticRegression(C=1/(X.shape[0]*lambda_reg), penalty='l1', solver='saga'),
    model_Y_final=WeightedLasso(alpha=lambda_reg)
)
```

```python
est.fit(Y, T, X=X, W=W)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:   26.6s
    [Parallel(n_jobs=-1)]: Done 176 tasks      | elapsed:   27.6s
    [Parallel(n_jobs=-1)]: Done 200 out of 200 | elapsed:   27.8s finished
    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    0.2s
    [Parallel(n_jobs=-1)]: Done 185 out of 200 | elapsed:    1.0s remaining:    0.0s
    [Parallel(n_jobs=-1)]: Done 200 out of 200 | elapsed:    1.0s finished

    <econml.orf._ortho_forest.DROrthoForest at 0x1a7b974ee48>

```python
# Calculate treatment effects for the default treatment points T0=0 and T1=1
treatment_effects = est.effect(X_test)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:   37.4s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:   41.0s finished

```python
# Calculate default (95%) confidence intervals for the default treatment points T0=0 and T1=1
te_lower, te_upper = est.effect_interval(X_test)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    1.8s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:    3.5s finished

```python
est2 = CausalForestDML(model_y=Lasso(alpha=lambda_reg),
                       model_t=LogisticRegression(C=1/(X.shape[0]*lambda_reg)),
                       n_estimators=200, min_samples_leaf=5,
                       max_depth=50, max_samples=subsample_ratio/2,
                       discrete_treatment=True,
                       random_state=123)
est2.fit(Y, T, X=X, W=W, cache_values=True)
treatment_effects2 = est2.effect(X_test)
te_lower2, te_upper2 = est2.effect_interval(X_test)
```

```python
est2.summary()
```

    Population summary of CATE predictions on Training Data

Uncertainty of Mean Point Estimate

  mean_point	 stderr_mean	  zstat	 pvalue	 ci_mean_lower	 ci_mean_upper

     3.088	      0.157	    19.677	   0.0	      2.78	          3.396

Distribution of Point Estimate

  std_point	 pct_point_lower	 pct_point_upper

    1.757	        0.846	           6.962

Total Variance of Point Estimate

  stderr_point	 ci_point_lower	 ci_point_upper

      1.764	         0.774	          6.951

Doubly Robust ATE on Training Data Results

   	   point_estimate	 stderr	  zstat	 pvalue	 ci_lower	 ci_upper

  ATE	      3.158	      0.082	 38.551	   0.0	    2.997	    3.318

Doubly Robust ATT(T=0) on Training Data Results

   	   point_estimate	 stderr	  zstat	 pvalue	 ci_lower	 ci_upper

  ATT	       3.1	       0.096	 32.322	   0.0	    2.912	    3.288

Doubly Robust ATT(T=1) on Training Data Results

   	   point_estimate	 stderr	  zstat	 pvalue	 ci_lower	 ci_upper

  ATT	      3.218	      0.134	 23.965	   0.0	    2.955	    3.481

Note: The stderr_mean is a conservative upper bound.

## 2.3. Performance Visualization

```python
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.title("DiscreteTreatmentOrthoForest")
plt.plot(X_test, treatment_effects, label='ORF estimate')
expected_te = np.array([exp_te(x_i) for x_i in X_test])
plt.plot(X_test[:, 0], expected_te, 'b--', label='True effect')
plt.fill_between(X_test[:, 0], te_lower, te_upper, label="95% BLB CI", alpha=0.3)
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.subplot(1, 2, 2)
plt.title("CausalForest")
plt.plot(X_test, treatment_effects2, label='ORF estimate')
expected_te = np.array([exp_te(x_i) for x_i in X_test])
plt.plot(X_test[:, 0], expected_te, 'b--', label='True effect')
plt.fill_between(X_test[:, 0], te_lower2, te_upper2, label="95% BLB CI", alpha=0.3)
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.show()
```

# 3. Example Usage with Multiple Treatment Synthetic Data

## 3.1. DGP
We use the following DGP:

\begin{align}
Y = & \sum_{t=1}^{n_{\text{treatments}}} 1\{T=t\}\cdot \theta_{T}(X) + \langle W, \gamma\rangle + \epsilon, \; \epsilon \sim \text{Unif}(-1, 1), \\
\text{Pr}[T=t \mid W] \propto & \exp\{\langle W, \beta_t \rangle\}, \;\;\;\; \forall t\in \{0, 1, \ldots, n_{\text{treatments}}\}
\end{align}

where $W$ is a matrix of high-dimensional confounders, $\beta_t, \gamma$ are sparse.

For this particular example DGP we used $n_{\text{treatments}}=3$ and
\begin{align}
\theta_1(x) = & \exp( 2 x_1 ),\\
\theta_2(x) = &  3 \cdot \sigma(100\cdot (x_1 - .5)),\\
\theta_3(x) = & -2 \cdot \sigma(100\cdot (x_1 - .25)),
\end{align}
where $\sigma$ is the sigmoid function.

```python
def get_test_train_data(n, n_w, support_size, n_x, te_func, n_treatments):
    # Outcome support
    support_Y = np.random.choice(range(n_w), size=support_size, replace=False)
    coefs_Y = np.random.uniform(0, 1, size=support_size)
    epsilon_sample = lambda n: np.random.uniform(-1, 1, size=n)
    # Treatment support
    support_T = support_Y
    coefs_T = np.random.uniform(0, 1, size=(support_size, n_treatments))
    eta_sample = lambda n: np.random.uniform(-1, 1, size=n)
    # Generate controls, covariates, treatments and outcomes
    W = np.random.normal(0, 1, size=(n, n_w))
    X = np.random.uniform(0, 1, size=(n, n_x))
    # Heterogeneous treatment effects
    TE = np.array([te_func(x_i, n_treatments) for x_i in X])
    log_odds = np.dot(W[:, support_T], coefs_T)
    T_sigmoid = np.exp(log_odds)
    T_sigmoid = T_sigmoid/np.sum(T_sigmoid, axis=1, keepdims=True)
    T = np.array([np.random.choice(n_treatments, p=p) for p in T_sigmoid])
    TE = np.concatenate((np.zeros((n,1)), TE), axis=1)
    Y = TE[np.arange(n), T] + np.dot(W[:, support_Y], coefs_Y) + epsilon_sample(n)
    X_test = np.array(list(product(np.arange(0, 1, 0.01), repeat=n_x)))

    return (Y, T, X, W), (X_test, np.array([te_func(x, n_treatments) for x in X_test]))
```

```python
import scipy.special
def te_func(x, n_treatments):
    return [np.exp(2*x[0]), 3*scipy.special.expit(100*(x[0] - .5)) - 1, -2*scipy.special.expit(100*(x[0] - .25))]

np.random.seed(123)
(Y, T, X, W), (X_test, te_test) = get_test_train_data(2000, 3, 3, 1, te_func, 4)
```

## 3.2. Train Estimator

```python
est = DROrthoForest(n_trees=500, model_Y = WeightedLasso(alpha=lambda_reg))
```

```python
est.fit(Y, T, X=X, W=W)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:   31.2s
    [Parallel(n_jobs=-1)]: Done 112 tasks      | elapsed:   32.9s
    [Parallel(n_jobs=-1)]: Done 272 tasks      | elapsed:   35.5s
    [Parallel(n_jobs=-1)]: Done 500 out of 500 | elapsed:   39.5s finished
    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    0.3s
    [Parallel(n_jobs=-1)]: Done 208 tasks      | elapsed:    3.4s
    [Parallel(n_jobs=-1)]: Done 500 out of 500 | elapsed:    7.9s finished

    <econml.orf._ortho_forest.DROrthoForest at 0x1a7bac95828>

```python
# Calculate marginal treatment effects
treatment_effects = est.const_marginal_effect(X_test)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:   20.6s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:   22.8s finished

```python
# Calculate default (95%) marginal confidence intervals for the test data
te_lower, te_upper = est.const_marginal_effect_interval(X_test)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    2.7s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:    5.1s finished

```python
res = est.const_marginal_effect_inference(X_test)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    2.5s
    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:    4.9s finished

```python
res.summary_frame()
```

      point_estimate
      stderr
      zstat
      pvalue
      ci_lower
      ci_upper

      X
      T

      0
      T0_1
      1.013
      0.159
      6.360
      0.000
      0.701
      1.325

      T0_2
      -0.989
      0.149
      -6.636
      0.000
      -1.281
      -0.697

      T0_3
      0.034
      0.226
      0.152
      0.879
      -0.408
      0.477

      1
      T0_1
      1.018
      0.160
      6.379
      0.000
      0.705
      1.331

      T0_2
      -0.987
      0.147
      -6.717
      0.000
      -1.276
      -0.699

      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...

      98
      T0_2
      1.967
      0.210
      9.345
      0.000
      1.554
      2.379

      T0_3
      -2.021
      0.163
      -12.414
      0.000
      -2.340
      -1.702

      99
      T0_1
      6.867
      0.244
      28.194
      0.000
      6.390
      7.344

      T0_2
      1.966
      0.212
      9.276
      0.000
      1.551
      2.382

      T0_3
      -2.017
      0.163
      -12.352
      0.000
      -2.337
      -1.697

300 rows × 6 columns

```python
est2 = CausalForestDML(model_y=Lasso(alpha=lambda_reg),
                       model_t=LogisticRegression(C=1/(X.shape[0]*lambda_reg)),
                       n_estimators=4000, min_samples_leaf=5,
                       max_depth=50, max_samples=subsample_ratio/2,
                       discrete_treatment=True,
                       random_state=123)
est2.fit(Y, T, X=X, W=W)
treatment_effects2 = est2.const_marginal_effect(X_test)
te_lower2, te_upper2 = est2.const_marginal_effect_interval(X_test, alpha=.01)
```

## 3.3. Performance Visualization

```python
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.title("DiscreteTreatmentOrthoForest")
y = treatment_effects
colors = ['b', 'r', 'g']
for it in range(y.shape[1]):
    plt.plot(X_test[:, 0], te_test[:, it], '--', label='True effect T={}'.format(it), color=colors[it])
    plt.fill_between(X_test[:, 0], te_lower[:, it], te_upper[:, it], alpha=0.3, color='C{}'.format(it))
    plt.plot(X_test, y[:, it], label='ORF estimate T={}'.format(it), color='C{}'.format(it))
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.subplot(1, 2, 2)
plt.title("CausalForest")
y = treatment_effects2
colors = ['b', 'r', 'g']
for it in range(y.shape[1]):
    plt.plot(X_test[:, 0], te_test[:, it], '--', label='True effect T={}'.format(it), color=colors[it])
    plt.fill_between(X_test[:, 0], te_lower2[:, it], te_upper2[:, it], alpha=0.3, color='C{}'.format(it))
    plt.plot(X_test, y[:, it], label='ORF estimate T={}'.format(it), color='C{}'.format(it))
plt.ylabel("Treatment Effect")
plt.xlabel("x")
plt.legend()
plt.show()
```

# 4. Example Usage with Real Continuous Treatment Observational Data

We applied our technique to Dominick’s dataset, a popular historical dataset of store-level orange juice prices and sales provided by University of Chicago Booth School of Business.

The dataset is comprised of a large number of covariates $W$, but researchers might only be interested in learning the elasticity of demand as a function of a few variables $x$ such
as income or education.

We applied the `ContinuousTreatmentOrthoForest` to estimate orange juice price elasticity
as a function of income, and our results, unveil the natural phenomenon that lower income consumers are more price-sensitive.

## 4.1. Data

```python
# A few more imports
import os
import pandas as pd
import urllib.request
from sklearn.preprocessing import StandardScaler
```

```python
# Import the data
file_name = "oj_large.csv"

if not os.path.isfile(file_name):
    print("Downloading file (this might take a few seconds)...")
    urllib.request.urlretrieve("https://msalicedatapublic.blob.core.windows.net/datasets/OrangeJuice/oj_large.csv", file_name)
oj_data = pd.read_csv(file_name)
oj_data.head()
```

      store
      brand
      week
      logmove
      feat
      price
      AGE60
      EDUC
      ETHNIC
      INCOME
      HHLARGE
      WORKWOM
      HVAL150
      SSTRDIST
      SSTRVOL
      CPDIST5
      CPWVOL5

      0
      2
      tropicana
      40
      9.018695
      0
      3.87
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927

      1
      2
      tropicana
      46
      8.723231
      0
      3.87
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927

      2
      2
      tropicana
      47
      8.253228
      0
      3.87
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927

      3
      2
      tropicana
      48
      8.987197
      0
      3.87
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927

      4
      2
      tropicana
      50
      9.093357
      0
      3.87
      0.232865
      0.248935
      0.11428
      10.553205
      0.103953
      0.303585
      0.463887
      2.110122
      1.142857
      1.92728
      0.376927

```python
# Prepare data
Y = oj_data['logmove'].values
T = np.log(oj_data["price"]).values
scaler = StandardScaler()
W1 = scaler.fit_transform(oj_data[[c for c in oj_data.columns if c not in ['price', 'logmove', 'brand', 'week', 'store']]].values)
W2 = pd.get_dummies(oj_data[['brand']]).values
W = np.concatenate([W1, W2], axis=1)
X = oj_data[['INCOME']].values
```

## 4.2. Train Estimator

```python
# Define some parameters
n_trees = 1000
min_leaf_size = 50
max_depth = 20
subsample_ratio = 0.04
```

```python
est = DMLOrthoForest(
        n_trees=n_trees, min_leaf_size=min_leaf_size, max_depth=max_depth,
        subsample_ratio=subsample_ratio,
        model_T=Lasso(alpha=0.1),
        model_Y=Lasso(alpha=0.1),
        model_T_final=WeightedLassoCVWrapper(cv=3),
        model_Y_final=WeightedLassoCVWrapper(cv=3)
       )
```

```python
est.fit(Y, T, X=X, W=W)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:   20.3s
    [Parallel(n_jobs=-1)]: Done 152 tasks      | elapsed:   21.0s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:   22.5s finished
    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    0.0s
    [Parallel(n_jobs=-1)]: Done 888 tasks      | elapsed:    1.6s
    [Parallel(n_jobs=-1)]: Done 1000 out of 1000 | elapsed:    2.1s finished

    <econml.orf._ortho_forest.DMLOrthoForest at 0x1a7cdd37588>

```python
min_income = 10.0
max_income = 11.1
delta = (max_income - min_income) / 100
X_test = np.arange(min_income, max_income + delta - 0.001, delta).reshape(-1, 1)
```

```python
# Calculate marginal treatment effects
treatment_effects = est.const_marginal_effect(X_test)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:   23.0s
    [Parallel(n_jobs=-1)]: Done 101 out of 101 | elapsed:   35.2s finished

```python
# Calculate default (95%) marginal confidence intervals for the test data
te_upper, te_lower = est.const_marginal_effect_interval(X_test)
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 8 concurrent workers.
    [Parallel(n_jobs=-1)]: Done  16 tasks      | elapsed:    6.1s
    [Parallel(n_jobs=-1)]: Done 101 out of 101 | elapsed:   21.3s finished

```python
est2 = CausalForestDML(model_y=WeightedLassoCVWrapper(cv=3),
                       model_t=WeightedLassoCVWrapper(cv=3),
                       n_estimators=n_trees, min_samples_leaf=min_leaf_size, max_depth=max_depth,
                       max_samples=subsample_ratio/2,
                       random_state=123)
est2.fit(Y, T, X=X, W=W)
treatment_effects2 = est2.effect(X_test)
te_lower2, te_upper2 = est2.effect_interval(X_test)
```

## 4.3. Performance Visualization

```python
# Plot Orange Juice elasticity as a function of income
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.plot(X_test.flatten(), treatment_effects, label="OJ Elasticity")
plt.fill_between(X_test.flatten(), te_lower, te_upper, label="95% BLB CI", alpha=0.3)
plt.xlabel(r'$\log$(Income)')
plt.ylabel('Orange Juice Elasticity')
plt.legend()
plt.title("Orange Juice Elasticity vs Income: ContinuousTreatmentOrthoForest")
plt.subplot(1, 2, 2)
plt.plot(X_test.flatten(), treatment_effects2, label="OJ Elasticity")
plt.fill_between(X_test.flatten(), te_lower2, te_upper2, label="95% BLB CI", alpha=0.3)
plt.xlabel(r'$\log$(Income)')
plt.ylabel('Orange Juice Elasticity')
plt.legend()
plt.title("Orange Juice Elasticity vs Income: CausalForest")
plt.show()
```

```python

```

notebooks/Causal Forest and Orthogonal Random Forest Examples.ipynb