Decoder

A decoder is a function that takes a latent variable and produces an output (Observations or Control inputs).

(model::Decoder)(x::AbstractArray, p::ComponentVector, st::NamedTuple)

The forward pass of the decoder.

Arguments:

• x: The input to the decoder.
• p: The parameters.
• st: The state.

returns:

- 'ŷ': The output of the decoder.
- 'st': The state of the decoder.

Encoder

An encoder is a container layer that contains three sub-layers: linear_net, init_net, and context_net.

Fields

• linear_net: A layer that maps the input to a hidden representation.
• init_net: A layer that maps the hidden representation to the initial hidden state.
• context_net: A layer that maps the hidden representation to the context.

(model::Encoder)(x::AbstractArray, p::ComponentVector, st::NamedTuple)

The forward pass of the encoder.

Arguments:

• x: The input to the encoder (e.g. observations).
• p: The parameters.
• st: The state of the encoder.

returns:

- `x̂₀`: The initial hidden state.
- `context`: The context.

HarmonicOscillators

A struct that defines a system of harmonic oscillators.

Arguments

• `N' : Number of oscillators
• `M' : Number of inputs
• `ω' : Frequency parameter ω
• `γ' : Damping coefficient γ
• `K' : Coupling matrix K
• `B' : External input coupling matrix B

(m::HarmonicOscillators)(x, u, t, p, st)

The forward pass of the HarmonicOscillators.

LatentUDE(obs_encoder, ctrl_encoder, dynamics, obs_decoder, ctrl_decoder)

Constructs a Latent Universal Differential Equation model.

Arguments:

• obs_encoder: A function that encodes the observations y to get the initial hidden state x₀ and context for the dynamics if needed (Partial observability)
• ctrl_encoder: A function that encodes (high-dimensional) inputs/controls to a lower-dimensional representation if needed.
• dynamics: A function that models the dynamics of the system (your ODE/SDE).
• obs_decoder: A function that decodes the hidden states x to the observations y.
• 'ctrl_decoder': A function that decodes the control representation to the original control space if needed.
• 'device': The device on which the model is stored. Default is cpu.

(model::LatentUDE)(y::AbstractArray, u::Union{Nothing, AbstractArray}, ts::AbstractArray, ps::ComponentArray, st::NamedTuple)

The forward pass of the LatentUDE model.

Arguments:

• y: Observations
• u: Control inputs
• ts: Time points
• ps: Parameters
• st: NamedTuple of states

Returns:

• ŷ: Decoded observations from the hidden states.
• ū: Decoded control inputs from the hidden states.
• x̂₀: Encoded initial hidden state.
• kl_path: KL divergence path. (Only for SDE dynamics, otherwise nothing)

Modern Wilson-Cowan Model of Neural Populations

The modern Wilson-Cowan model is a system of ordinary differential equations that describes the dynamics of excitatory and inhibitory neural populations with external inputs.

ODE(vector_field, solver; kwargs...)

Constructs an ODE model.

Arguments:

• vector_field: The vector field of the ODE.
• `solver': The nummerical solver used to solve the ODE.
• kwargs: Additional keyword arguments to pass to the solver.

(de::ODE)(x::AbstractArray, u::Union{Nothing, AbstractArray}, ts::AbstractArray, p::ComponentVector, st::NamedTuple)

The forward pass of the ODE.

Arguments:

• x: The initial hidden state.
• u: The control input.
• ts: The time steps.
• p: The parameters.
• st: The state.

returns: - The solution of the ODE. - The state of the model.

SDE(drift, drift_aug, diffusion, solver; kwargs...)

Constructs an SDE model.

Arguments:

• drift: The drift of the SDE.
• drift_aug: The augmented drift of the SDE. Used only for training.
• diffusion: The diffusion of the SDE.
• `solver': The nummerical solver used to solve the SDE.
• kwargs: Additional keyword arguments to pass to the solver.

(de::SDE)(x::AbstractArray, u::AbstractArray, c::AbstractArray, ts::StepRangeLen, p::ComponentVector, st::NamedTuple)

The forward pass of the SDE.

Arguments:

• x: The initial hidden state.
• u: The control input.
• c: The context.
• ts: The time steps.
• p: The parameters.
• st: The state.

returns: - The solution of the SDE. - The state of the model.

Wilson-Cowan Model of Neural Populations

The classic Wilson-Cowan model is a system of ordinary differential equations that describes the dynamics of excitatory and inhibitory neural populations.

Identity_Decoder()

Constructs an identity decoder. Useful for fully observable systems.

Identity_Encoder()

Constructs an identity encoder. Useful for fully observable systems.

Linear_Decoder(obs_dim, latent_dim)

Constructs a linear decoder.

Arguments:

• obs_dim: Dimension of the observations.
• latent_dim: Dimension of the latent space.
• noise: Type of observation noise. Default is Gaussian. Options are Gaussian, Poisson, None.

returns:

- The decoder.

MLP_Decoder(obs_dim, latent_dim, hidden_dim, n_hidden)

Constructs an MLP decoder.

Arguments:

• obs_dim: Dimension of the observations.
• latent_dim: Dimension of the latent space.
• hidden_dim: Dimension of the hidden layers.
• n_hidden: Number of hidden layers.
• noise: Type of observation noise. Default is Gaussian. Options are Gaussian, Poisson, None.

returns:

- The decoder.

Nothing_Decoder()

Constructs a decoder that does nothing.

Recurrent_Encoder(obs_dim, latent_dim, context_dim, hidden_dim, t_init)

Constructs a recurrent encoder.

Arguments:

• obs_dim: Dimension of the observations.
• latent_dim: Dimension of the latent space.
• context_dim: Dimension of the context.
• hidden_dim: Dimension of the hidden state.
• t_init: Number of initial time steps to use for the initial hidden state.

bits_per_spike(rates, spikes)

Compute the bits per spike by comparing the Poisson log-likelihood of the rates with the Poisson log-likelihood of the mean spikes.

Arguments:

• rates: The predicted rates.
• spikes: The observed spikes.

returns:

- The bits per spike.

forward!(model::LatentUDE, y::AbstractArray, u::Union{Nothing, AbstractArray}, ts::AbstractArray, ps::ComponentArray, st::NamedTuple, dynamics::ODE)

The forward pass of the LatentODE model.

forward!(model::LatentUDE, y::AbstractArray, u::Union{Nothing, AbstractArray}, ts::AbstractArray, ps::ComponentArray, st::NamedTuple, dynamics::SDE)

The forward pass of the LatentSDE model.

frange_cycle_linear(n_iter, start, stop, n_cycle, ratio)

Generate a linear schedule with cycles.

Arguments:

• n_iter: Number of iterations.
• start: Start value.
• stop: Stop value.
• n_cycle: Number of cycles.
• ratio: Ratio of the linear schedule.

returns:

- The linear schedule.

interp!(ts, cs, time_point)

Interpolates the control signal at a given time point.

Arguments:

• ts: Array of time points.
• cs: Array of control signals.
• time_point: The time point at which to interpolate the control signal.

returns:

- The interpolated control signal.

kl_normal(μ, σ²)

Compute the KL divergence between a normal distribution and a standard normal distribution.

Arguments:

• μ: Mean of the normal distribution.
• σ²: Variance of the normal distribution.

returns:

- The KL divergence.

mse(ŷ, y)

Compute the mean squared error.

Arguments:

• ŷ: Predicted values.
• y: Observed values.

returns:

- The mean squared error.

normal_loglikelihood(μ, σ², y)

Compute the log-likelihood of a normal distribution.

Arguments:

• μ: Mean of the normal distribution.
• σ²: Variance of the normal distribution.
• y: The observed values.

returns:

- The log-likelihood.

phaseplot(de::ODE, x₀_ranges, u, ts, p, st; kwargs...)

Plots the phase portrait of the ODE model.

Arguments:

• de: The ODE model.
• x₀_ranges: The initial condition ranges.
• u: The control input.
• ts: The time steps.
• p: The parameters.
• st: The state of the model.
• kwargs: Additional keyword arguments for plotting.

poisson_loglikelihood(λ, y)

Compute the log-likelihood of a Poisson distribution.

Arguments:

• λ: The rate of the Poisson distribution.
• y: The observed spikes.

returns:

- The log-likelihood.

predict(model::LatentUDE, y::AbstractArray, u::Union{Nothing, AbstractArray}, ts::AbstractArray, ps::ComponentArray, st::NamedTuple, n_samples::Int)

Samples trajectories from the LatentUDE model.

Arguments:

• model: The LatentUDE model to sample from.
• y: Observations used to encode the initial hidden state.
• u: Inputs for the input encoder. Can be Nothing or an array.
• ts: Array of time points at which to sample the trajectories.
• ps: Parameters for the model.
• st: NamedTuple of states for different components of the model.
• n_samples: Number of samples used to make the prediction.

Returns:

• ŷ: Decoded observations from the sampled hidden states * n_samples.
• ū: Decoded control inputs from the sampled hidden states * n_samples.
• x: Sampled hidden state trajectories * n_samples.

sample_dynamics(de::ODE, x̂₀, u, ts, p, st, n_samples)

Samples trajectories from the ODE model.

Arguments:

• de: The ODE model to sample from.
• x̂₀: The initial hidden state.
• u: Inputs for the input encoder. Can be Nothing or an array.
• ts: Array of time points at which to sample the trajectories.
• p: The parameters.
• st: The state.
• n_samples: The number of samples to generate.

returns: - The sampled trajectories. - The state of the model.

sample_dynamics(de::SDE, x̂₀, u, ts, p, st, n_samples)

Samples trajectories from the SDE model.

Arguments:

• de: The SDE model to sample from.
• x̂₀: The initial hidden state.
• u: Inputs for the input encoder. Can be Nothing or an array.
• ts: Array of time points at which to sample the trajectories.
• p: The parameters.
• st: The state.
• n_samples: The number of samples to generate.

returns: - The sampled trajectories. - The state of the model.

sample_rp(x::Tuple)

Samples from a MultiVariate Normal distribution using the reparameterization trick.

Arguments:

• x: Tuple of the mean and squared variance of a MultiVariate Normal distribution.

returns:

- The sampled value.