internals.gradients
Not for the user: the ODE run in JAX: gradients of results with respect to parameters, and many runs at once
A compartmental engine flattens its flows into a ss.kernels.Program. The kernels (Numba or Rust) step through it; run() here takes the same program and runs the same fourth-order Runge-Kutta (or Euler) steps in JAX, so that the run can be differentiated (gradient()) and many runs with different values can be taken in one call (run_sets()). JAX is imported only when one of these is used.
What JAX differentiates is the run: the program’s numbers, and the compartments at the start, to the compartments and flows at every step. The steps either side of it are Python (working out those numbers from the parameters, and the results from the run, including outcomes and the goodness of fit), and are differentiated by central differences along the direction the run is differentiated in. They are cheap and smooth, so this is accurate to about 1e-9 of the value.
Used with: compartments; declared modules
Classes
| Name | Description |
|---|---|
| Density | The log of the posterior density of a sim’s fitted parameters, as one function in JAX, so that it can be compiled and differentiated whole (see sim.fit() with method=‘nuts’) |
| JointRun | The diseases of joint states, run together by the joint states’ program (see ss.solvers.JointStates) |
| LogDensity | The log of the posterior density of a sim’s fitted parameters as a function of their values, in JAX, for use outside sim.fit(): see ss.Sim.log_density() |
| Run | One engine’s run, as a function of the numbers in it |
| Unsupported | Something in the model that the log-density in JAX does not cover yet |
Density
internals.gradients.Density(
sim,
fitted,
x0,
set_pars,
likelihood=None,
weights=None,
external=None,
bounds=None,
)The log of the posterior density of a sim’s fitted parameters, as one function in JAX, so that it can be compiled and differentiated whole (see sim.fit() with method=‘nuts’)
The numbers in the run are worked out from the parameters in JAX: the values in the rates by evaluating the model’s own expressions, and the rest (e.g. the compartments at the start) from the change in them as each parameter changes, which has to be linear (checked at the corners and the middle of the bounds when it is built). Then the run, the results that have data, and the likelihood and priors. If anything is not covered, Unsupported says what.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| sim | ss.Sim |
an initialized sim that ss.solvers.why_not(sim, 'jax') allows |
required |
| fitted | dict |
the fitted parameters, by name | required |
| x0 | array |
their values to start from | required |
| set_pars | func |
restarts the sim with an array of values | required |
| likelihood | dict |
the likelihood of each data series, as for ss.fitting.loglik() |
None |
| weights | dict |
a multiplier on each series’ log-likelihood | None |
| external | list |
for each parameter, whether it belongs to a likelihood given at calibration time (rather than to the model) | None |
| bounds | tuple |
the lowest and highest value of each parameter, over which a step is allowed as many substeps as the run needs (default: each parameter’s bounds) |
None |
Methods
| Name | Description |
|---|---|
| capture | The numbers in the run of each engine with some parameter values, worked out in Python |
| check_linear | Check that the numbers worked out from their slopes at the start (see numbers()) are what they are at other points (the |
| convert | A fitted value as the engine uses it: in the sim’s unit (see convert_fitted()) |
| corners | The points where the substeps a step needs are worked out: the corners of the bounds (or with more than 4 parameters, each parameter at each of its bounds, the others at their initial values) |
| describe | What a number in the run of engine j is, for a message: the entry where of key |
| evaluate | The log of the posterior density at some parameter values, and the most substeps a step of any run needed (more than it took, if clamped by n_max) |
| init_outputs | Which result each data series is compared with, and the matrix that lines it up with the data |
| init_program | Work out which numbers in the run depend on which parameters, and how |
| init_several | For each parameter with several values (see ss.fitting.expand()), how its value is worked out from the fitted ones: a value |
| limits | The lowest and highest value of each parameter (its initial value for one without bounds) |
| loglik | As ss.fitting.logpdf(), in JAX, with the fitted parameters of the likelihood at their values in x |
| numbers | The numbers in the run of each engine, as JAX arrays, at some parameter values |
| outcomes | The outcomes, as ss.solvers.Compartmental.derived_outcomes() works them out |
| results | The results that have data, from the run of each engine: a dict by data column |
| reverse | Make the runs differentiable in reverse mode (each step takes its most substeps), for jax.grad; forward mode, as NUTS uses, is faster without |
| states | The number in each state at each time point, at some parameter values: a dict of JAX arrays by ‘module.state’ |
| values | The values of a parameter with several values, in the sim’s unit: those fitted from x, the others as given |
capture
internals.gradients.Density.capture(x)The numbers in the run of each engine with some parameter values, worked out in Python
check_linear
internals.gradients.Density.check_linear(checked, rtol=1e-06, atol=1e-09)Check that the numbers worked out from their slopes at the start (see numbers()) are what they are at other points (the corners and the middle of the bounds), and raise Unsupported naming those that are not, e.g. a fraction of q*q
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| checked | list |
each point, and the numbers in the run of each engine there | required |
| rtol | float |
how far a number may be from its value, relative to it (the slopes, by central differences, put a number that is linear off by about 1e-12 of it, or 1e-8 with bounds 10^4 times the starting value) | 1e-06 |
| atol | float |
the same, relative to the largest of its kind (e.g. for a number found as the difference of two large ones) | 1e-09 |
convert
internals.gradients.Density.convert(i, val)A fitted value as the engine uses it: in the sim’s unit (see convert_fitted())
corners
internals.gradients.Density.corners()The points where the substeps a step needs are worked out: the corners of the bounds (or with more than 4 parameters, each parameter at each of its bounds, the others at their initial values)
describe
internals.gradients.Density.describe(j, key, where)What a number in the run of engine j is, for a message: the entry where of key
evaluate
internals.gradients.Density.evaluate(x)The log of the posterior density at some parameter values, and the most substeps a step of any run needed (more than it took, if clamped by n_max)
init_outputs
internals.gradients.Density.init_outputs()Which result each data series is compared with, and the matrix that lines it up with the data
init_program
internals.gradients.Density.init_program()Work out which numbers in the run depend on which parameters, and how
init_several
internals.gradients.Density.init_several()For each parameter with several values (see ss.fitting.expand()), how its value is worked out from the fitted ones: a value for each level as each group’s (levels, the matrix from the levels to the groups), and a time series by interpolating them (series_of), both exactly, as a parameter with one value is; an array (arrays) as the numbers in the run that change linearly with it. Also the value of each entry as the parameter holds it (raws), for convert().
limits
internals.gradients.Density.limits()The lowest and highest value of each parameter (its initial value for one without bounds)
loglik
internals.gradients.Density.loglik(dist, observed, model, x)As ss.fitting.logpdf(), in JAX, with the fitted parameters of the likelihood at their values in x
numbers
internals.gradients.Density.numbers(x)The numbers in the run of each engine, as JAX arrays, at some parameter values
outcomes
internals.gradients.Density.outcomes(engine, ctx, new, npts)The outcomes, as ss.solvers.Compartmental.derived_outcomes() works them out
results
internals.gradients.Density.results(x, stocks, flows)The results that have data, from the run of each engine: a dict by data column
reverse
internals.gradients.Density.reverse()Make the runs differentiable in reverse mode (each step takes its most substeps), for jax.grad; forward mode, as NUTS uses, is faster without
states
internals.gradients.Density.states(x)The number in each state at each time point, at some parameter values: a dict of JAX arrays by ‘module.state’
values
internals.gradients.Density.values(positions, x, engine)The values of a parameter with several values, in the sim’s unit: those fitted from x, the others as given
JointRun
internals.gradients.JointRun(joint, run)The diseases of joint states, run together by the joint states’ program (see ss.solvers.JointStates)
Called with a tuple of the numbers of each disease’s run: they are put together into the numbers of the program (ss.solvers.JointStates.numbers()), which is run as any other, and its output is split between the diseases again. Returns what Run.function() does, with a tuple over the diseases for each of the first three.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| joint | ss.solvers.JointStates |
the joint states | required |
| run | Run | the run of its program | required |
LogDensity
internals.gradients.LogDensity(sim, pars=None, likelihood=None)The log of the posterior density of a sim’s fitted parameters as a function of their values, in JAX, for use outside sim.fit(): see ss.Sim.log_density()
Attributes
| Name | Type | Description |
|---|---|---|
| names | list |
the parameters, in the order of the values |
| x0 | array |
their current values |
| fn | func |
the function itself, of a JAX array of values, which can be composed with JAX (jax.grad, jax.vmap, …) |
Methods
| Name | Description |
|---|---|
| apply | One of the compiled functions, at some values (compiled the first time) |
| grad | Its gradient, by reverse-mode automatic differentiation (one pass, whatever the number of parameters) |
| run | The number in each state at each time point, at some values: a dict of arrays by ‘module.state’ |
apply
internals.gradients.LogDensity.apply(func, x)One of the compiled functions, at some values (compiled the first time)
grad
internals.gradients.LogDensity.grad(x)Its gradient, by reverse-mode automatic differentiation (one pass, whatever the number of parameters)
run
internals.gradients.LogDensity.run(x)The number in each state at each time point, at some values: a dict of arrays by ‘module.state’
Run
internals.gradients.Run(program, dt, npts, solver, n_sub, n_max)One engine’s run, as a function of the numbers in it
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| program | Program |
the engine’s program, for how its flows are put together | required |
| dt | float |
the timestep | required |
| npts | int |
the number of steps | required |
| solver | int |
0 for RK4, 1 for Euler (see ss.kernels.solvers) |
required |
| n_sub | int |
the number of substeps a step, or 0 for enough that no compartment turns over, or grows through its inflows, by more than a quarter in one | required |
| n_max | int |
the most substeps a step can take; a run that needs more is taken again with more | required |
Methods
| Name | Description |
|---|---|
| checkpoint | Whether a run differentiated in reverse mode keeps the state at only about √T of its T steps (see checkpointed()): if what reverse mode would otherwise keep (the compartments at each substep and stage) is large |
| comp_matrices | For nc compartments: the compartments by flows (-1 where a flow leaves) and by targets (1 where a target not mapped between groups takes people) |
| deliver | The deliveries at time point ti, at the start of the step from it, in order (see ss.kernels.apply_shifts()): each moves some of the people in each group to another, or (a move between states) into another state |
| derivs | The rate of change of each compartment, and of each transition and mover |
| evaluate | The value of the operations a to b-1 of the program, by group, with the compartments at ys (as ss.kernels.step() works out a rate) |
| frac | The time fraction between places times values by place and level (as ss.kernels.frac()) |
| function | The run, as a function of a dict of its numbers (see capture()): returns the compartments after the last step, at the start of each step, the flows over each step, and the most substeps a step needed |
| get | The compiled run: ‘run’ for one, ‘many’ for a batch, ‘jvp’ for its derivative along a batch of directions |
| key | What a compiled run depends on, apart from the numbers in it |
| layout | Work out, once, the matrices that put the flows together |
| members | The compartments of set s, by group, each times its weight if it has one (as ss.kernels.member()) |
| prepare | The numbers of a run, with what is worked out from them once a run: the time fraction as a matrix from place to |
| rates | The rate of every flow, by group, with the compartments at ys and the program’s values in vals: as ss.kernels.step() works them out |
| stepper | One step of the run, as a function of the compartments at its start and its number: returns the compartments at its end, the flows over it, and the substeps it needed |
checkpoint
internals.gradients.Run.checkpoint(X)Whether a run differentiated in reverse mode keeps the state at only about √T of its T steps (see checkpointed()): if what reverse mode would otherwise keep (the compartments at each substep and stage) is large
comp_matrices
internals.gradients.Run.comp_matrices(nc)For nc compartments: the compartments by flows (-1 where a flow leaves) and by targets (1 where a target not mapped between groups takes people)
deliver
internals.gradients.Run.deliver(y, X, ti)The deliveries at time point ti, at the start of the step from it, in order (see ss.kernels.apply_shifts()): each moves some of the people in each group to another, or (a move between states) into another state
derivs
internals.gradients.Run.derivs(ys, rates, X, h)The rate of change of each compartment, and of each transition and mover
evaluate
internals.gradients.Run.evaluate(ys, a, b, consts, exposures, local_n, regs)The value of the operations a to b-1 of the program, by group, with the compartments at ys (as ss.kernels.step() works out a rate)
frac
internals.gradients.Run.frac(v, block, X, by_columns)The time fraction between places times values by place and level (as ss.kernels.frac())
function
internals.gradients.Run.function(clamp=False)The run, as a function of a dict of its numbers (see capture()): returns the compartments after the last step, at the start of each step, the flows over each step, and the most substeps a step needed
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| clamp | bool |
see stepper() |
False |
get
internals.gradients.Run.get(kind='run')The compiled run: ‘run’ for one, ‘many’ for a batch, ‘jvp’ for its derivative along a batch of directions
key
internals.gradients.Run.key(program, dt, npts, solver, n_sub)What a compiled run depends on, apart from the numbers in it
layout
internals.gradients.Run.layout()Work out, once, the matrices that put the flows together
members
internals.gradients.Run.members(ys, s, consts)The compartments of set s, by group, each times its weight if it has one (as ss.kernels.member())
prepare
internals.gradients.Run.prepare(X)The numbers of a run, with what is worked out from them once a run: the time fraction as a matrix from place to place, if it is dense (a product with it is much faster in XLA than adding into an array at indices: 4x for 50 places), and its transpose (since XLA transposes a matrix that is an input again at every step: 10x slower for 260 places)
rates
internals.gradients.Run.rates(ys, vals)The rate of every flow, by group, with the compartments at ys and the program’s values in vals: as ss.kernels.step() works them out
stepper
internals.gradients.Run.stepper(X, clamp=False)One step of the run, as a function of the compartments at its start and its number: returns the compartments at its end, the flows over it, and the substeps it needed
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| X | dict |
the numbers in the run (see capture()) |
required |
| clamp | bool |
if True, a step takes at most n_max substeps, so that the run is the same function everywhere (for sampling, where it cannot be taken again with more) |
False |
Unsupported
internals.gradients.Unsupported()Something in the model that the log-density in JAX does not cover yet
It exists so that the reason can be given to the user rather than a JAX error.
Functions
| Name | Description |
|---|---|
| gradient | The value of something read from a sim’s run, and its derivative with respect to each of a set of values |
| jax | Import JAX (imported only when first used, since it takes a few tenths of a second), on the platform that ss.options.jax says |
| of_results | The derivative of results, or of a function of the run sim, with respect to parameters: see ss.Sim.gradient() |
| run | Take a whole run of an engine in JAX: returns what ss.kernels.run() does |
| run_sets | Run a sim with many sets of parameter values, taking the runs in one call to JAX |
gradient
internals.gradients.gradient(sim, set_pars, x0, read, rel_step=1e-05)The value of something read from a sim’s run, and its derivative with respect to each of a set of values
The run is differentiated by JAX (forward mode, one direction per value); see the module docstring for the rest.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| sim | ss.Sim |
an initialized sim that ss.solvers.why_not(sim, 'jax') allows, which is left as its run at x0 |
required |
| set_pars | func |
called with an array of values; restarts the sim with them (e.g. lambda x: sim.restart(pars=dict(beta=x[0]))) |
required |
| x0 | array |
the values to work out the derivative at | required |
| read | func |
called with the sim after its run; returns a number or an array (or an sc.obj of arrays) |
required |
| rel_step | float |
the step for the central differences, relative to each value (or absolute, for a value of 0) | 1e-05 |
Returns
| Name | Type | Description |
|---|---|---|
value, derivative: what read() returned at x0, and its derivative with respect to each value (a list, one per value, of the same type as the value) |
jax
internals.gradients.jax()Import JAX (imported only when first used, since it takes a few tenths of a second), on the platform that ss.options.jax says
If JAX has not been imported yet, it is given a device for each core of the CPU (as numpyro.set_host_device_count() does), or as many as the threads asked for in the environment (e.g. OMP_NUM_THREADS=1: see ss.settings.thread_limit), so that the chains of a posterior run on cores of their own; everything else runs on the first. XLA’s own threads (most of them compiling) are not limited by it: no flag of XLA’s limits them.
The platform it runs on is ss.options.jax (e.g. “cpu”), if not “auto”. If JAX has started already on others (by Starsim or anything else), it is started again on these: arrays made already stay on their devices, and are moved when used with new ones, and what was compiled is compiled again. “auto” goes back to the platforms JAX had before Starsim first changed them. The first time, its devices are started here, so that a GPU that fails to start (e.g. with a CUDA plugin of another version than jaxlib) says what to do.
of_results
internals.gradients.of_results(sim, results=None, pars=None)The derivative of results, or of a function of the run sim, with respect to parameters: see ss.Sim.gradient()
run
internals.gradients.run(engine)Take a whole run of an engine in JAX: returns what ss.kernels.run() does
run_sets
internals.gradients.run_sets(sim, par_sets)Run a sim with many sets of parameter values, taking the runs in one call to JAX
Each set is a dict of values by name, as for sim.restart(pars=...). Only for a sim that ss.solvers.why_not(sim, 'jax') allows: an ODE that the kernels run from start to finish.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| sim | ss.Sim |
an initialized sim | required |
| par_sets | list |
the values of each run | required |
Returns, for each declared module, its results over the runs, as ss.solvers.Compartmental.run_many() does.