*! version 1.0 2021-02-06

Title

mlad -- Maximum likelihood estimation with automatic differentiation using Python.

Syntax

mlad
eq
[eq ...]
[if] [in],
llfile(filename)
[model_options]

adtype(adtype) type of automatic differentiation to calculate hessian

id(varname) name of id variable to pass to Python

init(init) set initial values

llfile(filename) filename of Python likelihood file

nojit Do not invoke just in time compliation

matrices(mat list) list of matrices to pass to Python

matnames(list) names of matrices in Python

mlmethod(method) method of estimation

othervars(varlist) variables to be passed to Python

othervarsnames(list) names of variables in Python

pygradient python gradient function provided

pyhessian python Hessian function provided

othervarsnames(list) names of variables in Python

pysetup(filename) python file for data manipulation prior to estimation

robustok calculate robust standard errors

scalars(scalar list) list of scalars to pass to Python

scalarnames(list) names of scalars in Python

search(search option) initial values search options

staticscalars(scalar list) scalars treated as static arguments in Python

{[R] maximize} control the maximization process; seldom used

Description

mlad maximizes a log-likelihood function where the likelihood function is programmed in Python.
This enables the gradients and Hessian matrix to be obtained using automatic differentiation and to take advantage of using multiple CPUs.
With large datasets mlad tends to be substantially faster than ml
and has the important advantage that you don't have to derive the gradients
and the Hessian matrix analytically as these are obtained using automatic differentiation.

mlad uses Stata's inbuilt optimizer, ml, but makes a call to Python to calculate the
log-likelihood function, gradients and Hessian matrix. This means that the likelihood function must be written as a Python function. The functions to derive the gradients and the Hessian matrix are obtained using automatic differentiation.

mlad makes use the the Jax Python library.
The Python function for the log-likelihood needs to use Jax's versions of the numpy and/or scipy Python libraries.
Some of the advantages of using Jax are as follows.

1) Jax incorprates automatic differentiation. This means thee is no need and sit and derive the gradient and hessian functions for your log-likelihood.

2) Jax uses a XLA compiler for the Python code. This allows it to use multiple CPUs and potentially GPUs. This makes computation fast and allows those with multple CPUs to get the benefit, even without Stata MP.

3) Jax can use just-in-time compilation to further increase computational efficiency.

4) Jax has has automatic vectorization, to improve performance - useful, for example, for numerical integration.

In order to use mlad you need to have Stata 16 or above and access to Python.
You also need certain Python modules installed.
These are jax, jaxlib, numpy, scipy and importlib.
Currently, there is not possible to directly install the jaxlib module for Windows, but it is possible to
build from source.
I can't help with this process.

Factor variables: mlad works with factor variables. However if speed is important to you, then you should create indicator variables / interactions yourself as this will be faster.

Some examples of using mlad can be found here, https://pclambert.net/software/mlad.

Options

adtype(adtype) the type(s) of automatic differentiation used when calculation the Hessian function.
This can be any combination of forward (fwd) and reverse (rev),
i.e. revrev, revfwd, fwdrev, fwdfwd.
The default is revrev.

id(varname) gives the name of an id variable. This will be required if the log-likelihood needs to
calculated over groups, for example in random effects models. If the robustok option is used then
cluster robust standard errors will be calculated.

init(varname) sets the initial values, b_0. See ml

nojit does not invoke just-in-time compilation.
This can be useful when debugging as Python error messages can be more informative.
In addition, to be jitable Python functions have to be written in a particular style
and so it may not be possible to always use jit.
When writing the Python likelihood function, it is useful to first use the nojit option,
get the function working and then check that
it works without the nojit option.
See https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html for advice on making functions "jitable".

llfile(filename) python filename containing log-likelhood function. The function must be named python_ll.

matrices(matrix list) will pass matrices to Python if required for estimation of the log-likelihood. For example,the location of the knots when incorporating restricted cubic splines.
These will be stored in a dictionary in Python with either the same names specified in matrix list or
the names specified in matnames() option.

matnames(list) list of names that the matrices will be named in Python.
If not specified these will default to the same names in the matrices option.
However, these will often be ugly, tempory names and so it useful to give them more meaningful names.

mlmethod(method) asks to use a ml method other than the default of d2. Only dtype estimators are available.
This can be useful to check whether the gradients (d1debug) or Hessian (d2debug) have been calculated correctly
or for calculating the speed improvments when using automatic differentiation.

othervars(varlist) give the names of other variables required for the estimation of the log-likelihood.
This will include the outcome variable(s).
Weights and offsets are automatically passed to the Python likelhood function.
You do not need to include covariates in this option as these are automatically passed to the log-likelhood function.
These will be stored in a dictionary in Python with either the same names specified in varlist or the names specified in othervarnames() option.

othervarnames(list) list of names that the variables will be named in Python.
If not specified these will default to the same names in the othervars option.
However, these will often be ugly tempory names and so it useful to give them more meaningful names.

pygradient Python gradient function is supplied in llfile.
This means that automatic differentiation will not be used to calculate the gradient vector.
The function must be named python_grad with the same options as the log likelihood function.

pyhessian Python Hessian function is supplied in llfile.
This means that automatic differentiation will not be used to calculate the Hessian matrix.
The function must be named python_hessian with the same options as the log likelihood function.

pysetup(filename) Python file to execute prior to estimation.

robustok Will calculate robust stand errors.
Note normally robust standard errors cannnot be obtained with d-type estimators, but mlad uses automatic differentiation to derive the score equations.
If the id option has been specified these will be cluster robust standard errors.

scalars(scalar list) will pass scalars to Python if required for estimation of the log-likelihood.
These will be stored in a dictionary in Python with either the same names specified in scalar list or the names specified in scalarnames() option.
Note that it be necessary to pass some scalars as static arguments so that the
Python likelhood function is "jitable" - see the staticscalars() option.

scalarnames(list) list of names that the scalars will be named in Python.
If not specified these will default to the same names in the scalars option.
However, these will often be ugly, tempory names and so it useful to give them more meaningful names.

search(search option) gives the search option to be paseed to ml. See ml

staticscalars(scalar list) will pass scalars to Python if required for estimation of the log-likelihood that are treated
as static arguments when using just in time compilation.
This is necessary for arguments that affect the size of Python arrays.
For example, the number of groups in a random effects model.

Writing the Python log-likelhood function

mlad requires a Python function to calculate the log-likelhood.
Once this is written and passed via the llfile() option to mlad the fitted model will be a Stata d2 estimate and thus various postestimation commands can be used.

There are some important rules when writing the Python program.

1) The function needs to be named python_ll()

2) The first 4 arguements of the Python function will always be def python_ll(beta,X,wt,M):

beta - is Python list containing a vector of parameters for each equation.

X - is a Python list containing an array of covariates for each equation.

wt - is a vector of weights or a column of 1's of not specified.

M - is a Python dictionary containing addition varibles, matrices or scalars passed using the othervars(), matrices() or scalars() options.
They can be referenced in Python using the same names as they were named in Stata or the names given in the othervarames(), matnames() or scalarnames() options.
In addition if the pysetup() option is used additional information can be stored in the Python dictionary.

There will be only additional function arguments if the staticscalars() option is specified.

3) Jax modules should be imported, for example rather than using import numpy as np you should use import jax.numpy as jnp. Importing other modules, such as pandas will very likely result in an error.

4) The function should be "pure".
This means that all input data is passed through the function parameters (so you can't use Python global objects) and all results are passed through function results.
For the purposes of mlad the only thing that should be returned is the log-likelihood as a single scalar.

5) There are various other issues that can catch you out. See JAX The Sharp Bits for discussion of style issues.

To show some of these issues a simple example follows,

We will maximize a Weibull regression model, i.e. the same as using streg with the dist(weibull) option. This is a survival model. The survival function for a Weibull model is,

S(t) = exp(-lambda*t^gamma)

with hazard function

h(t) = lambda*gamma*t^(gamma - 1)

Both lambda and gamma must be positive thus when estimating the parameters it is usual to model log(lambda) and log(gamma).
Incoporating covariates for log(lambda) means that the estimated coefficents are log hazard ratios if log(gamma) does not vary by covariates.

With survival data the log-likelhood conribution of the ith individual can be expressed as,

lli = d_i*ln[h(t_i)] + log[S(t_i)]

Where t_i is the survival time and d_i the event indicator for the ith individual. So for the Weibull model
the log-likelhood can be expressed as,

lli = d_i*(ln(lambda) + ln(gamma) + (gamma-1)*ln(t)) - lambda*t^(gamma-1)

If using ml then an ado file would be written, here named weib_d0. Using a d0 evaluator, this could be,

program define weib_d0
version 16.1
args todo b lnf g H

tempvar lnlambda lngamma
mleval `lnlambda' = `b', eq(1)
mleval `lngamma' = `b', eq(2)

mlsum `lnf' = _d*(`lnlambda' + `lngamma' + (exp(`lngamma') - 1)*ln(_t)) - ///
exp(`lnlambda')*_t^(exp(`lngamma'))
if (`todo'==0 | `lnf'>=.) exit
end

For details, see ml, but in brief, the program extracts the linear predictor for both ln(lambda) and ln(gamma) using mleval and then feeds these into the log-likelhood function, which is summed using mlsum.

An example of fiting the model can be seen below

. webuse brcancer
. stset rectime, failure(censrec==1) scale(365.24)
. rename x1 age
. ml model d0 weib_d0 (ln_lambda: = hormon age) (ln_gamma: = hormon), maximize
. ml display

This will give the same parameter estimates as streg.

Fitting the same model using mlad requires writing a Python program to calculate the likelihood. This is shown below.

import jax.numpy as jnp
import mladutil as mu

def python_ll(beta1,X,wt,M):
lnlam = mu.linpred(beta,X,1)
lngam = mu.linpred(beta ,X,2)
gam = jnp.exp(lngam)

return(jnp.sum(d*(lnlam + lngam + (gam - 1)*jnp.log(t)) - jnp.exp(lnlam)*t**(gam)))

First two modules are imported. The first is JAX's version of numpy.
This will nearly always have to be imported.

The second, mladutil, is a set of utility programs for mlad.

The function name must always be python_ll.
There are 4 function arguments (as statics scalars are not needed for this example).

The first argument, beta, is a list with the first item the parameters for ln(lambda) and the second item the parameters for ln(gamma).

The second function arguement is X. The covariates are automatically transferred to Python and stored in a list with the covariates for the first equation in X[1] and the kth equation in X[k].
If any offsets have been specified, these will also be included in X[0].

The third argument defines any weights that have been specified or a columns of 1's if they have not been specified.

The final arguement, M, is a dictionary containing any variables specified in the othervars() option of mlad,
matrices specified in the matrices() option or
scalars specified in the scalars() option.
Here the survival time (_t) and the event indicator (_d) are needed to calculate the liklehood function.
Note that these will be named t and d in the Python dictionary, M, as defined in the othervarnames() option. See below.

linpred is a utility function to get the current predicted value for the kth equation given X and beta.
It is recommended that you use this function. linpred will automatically incorporate any offsets if specified.

The likelihood is the same as that specified in the Stata ado file. Note that this needs to be summed over observations and thus is returned as a single scalar.

The syntax for mlad in terms of specifiying equations is the same as ml.

. mlad (ln_lambda: = hormon age) ///
(ln_gamma: = hormon) ///
othervars(_t _d) ///
othervarnames(t d) ///
llfile(weib_like_jax)
. ml display

The are three additional options of mlad.
The othervars() option specifies that _t and _d will be passed to the
Python likelihood function in a Python dictionary.
The othervarsnames() option defines the keys used in the Python dictionary.
If not specified these woudl be the same as specified in the othervars() option.
Note that mlad always calls ml in non-interactive mode, so there is no maximize option.

Description

Some examples of using mlad can be found here, https://pclambert.net/software/mlad.