pseudodynamics.models._pde_informed_params¶

Classes

`log_pde_params`(channels[, growth_weight, ...])
`logrithmic_pde`(channels, step_size[, ...])
`pde_neighborloss`(channels, n_grid[, ...])
`pde_params`(channels[, growth_weight, ...])	Default model : PINN prediction + NeuralODE simulation to estimate parameters
`pde_params_base`(channels[, collapse_D, ...])
`pde_params_fastmode`(channels[, ...])
`pde_params_meshgrid`(channels, n_grid[, ...])	mlp g,v,D for meshgrid dataset
`pde_singlebranch_twotimepoints`([n_grid, ...])
`pde_u_free`(channels, step_size[, ...])

class pseudodynamics.models._pde_informed_params.log_pde_params(channels, growth_weight=None, collapse_D=True, collapse_v=False, g_channels=None, v_channels=None, D_channels=None, time_sensitive=True, lr=0.0003, ode_tol=0.0001, activation_fn='Tanh', deltax_weight=None, D_penalty=None, weight_intensity=None, time_scale_factor=None, pop_weight=None)[source]¶

Bases: pde_params

Parameters:: activation_fn (str | list)

equation(s, t)[source]¶

the log Reaction-Advection Diffusion equation

Return type:: tuple

class pseudodynamics.models._pde_informed_params.logrithmic_pde(channels, step_size, collapse_D=True, collapse_v=False, g_channels=None, v_channels=None, D_channels=None, time_sensitive=True, lr=0.0003, ode_tol=0.0001, activation_fn='Tanh', deltax_weight=None, D_penalty=None, weight_intensity=None, time_scale_factor=None)[source]¶

Bases: pde_u_free

Parameters:: activation_fn (str | list)

equation(s, t, logu, dloguds)[source]¶

Apply torch’s auto grad to compute the dynamics

Return type:: tuple

based on the following equation:: ∂u/∂t = ∂/∂s[ D* ∂u/∂s ] - ∂/∂s[ v*u ] + g*u

we calcuate the left hand side (lhs) and the right hand side

class pseudodynamics.models._pde_informed_params.pde_neighborloss(channels, n_grid, collapse_D=True, collapse_v=False, g_channels=None, v_channels=None, D_channels=None, time_sensitive=True, lr=0.0003, ode_tol=0.0001, activation_fn='Tanh', D_penalty=None, weight_intensity=None)[source]¶

Bases: pde_params_meshgrid

Parameters:: activation_fn (str | list)

forward(t, states)[source]¶

Same as torch.nn.Module.forward().

Args:: *args: Whatever you decide to pass into the forward method. **kwargs: Keyword arguments are also possible.
Return:: Your model’s output

training_step(train_batch, index)[source]¶

Here you compute and return the training loss and some additional metrics for e.g. the progress bar or logger.

Args:

batch: The output of your data iterable, normally a DataLoader. batch_idx: The index of this batch. dataloader_idx: The index of the dataloader that produced this batch.

(only if multiple dataloaders used)

Return:

Tensor - The loss tensor
dict - A dictionary which can include any keys, but must include the key 'loss' in the case of automatic optimization.
None - In automatic optimization, this will skip to the next batch (but is not supported for multi-GPU, TPU, or DeepSpeed). For manual optimization, this has no special meaning, as returning the loss is not required.

In this step you’d normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something model specific.

Example:

def training_step(self, batch, batch_idx):
    x, y, z = batch
    out = self.encoder(x)
    loss = self.loss(out, x)
    return loss

To use multiple optimizers, you can switch to ‘manual optimization’ and control their stepping:

def __init__(self):
    super().__init__()
    self.automatic_optimization = False


# Multiple optimizers (e.g.: GANs)
def training_step(self, batch, batch_idx):
    opt1, opt2 = self.optimizers()

    # do training_step with encoder
    ...
    opt1.step()
    # do training_step with decoder
    ...
    opt2.step()

Note:: When accumulate_grad_batches > 1, the loss returned here will be automatically normalized by accumulate_grad_batches internally.

class pseudodynamics.models._pde_informed_params.pde_params(channels, growth_weight=None, collapse_D=True, collapse_v=False, g_channels=None, v_channels=None, D_channels=None, time_sensitive=True, lr=0.0003, ode_tol=0.0001, activation_fn='Tanh', R_weight=None, deltax_weight=None, D_penalty=None, weight_intensity=None, time_scale_factor=None, pop_weight=None, cfm_weight=None, D_var_weight=None, neuralode_weight=None)[source]¶

Bases: pde_params_base

Default model : PINN prediction + NeuralODE simulation to estimate parameters

Parameters:

channel (list) – the number of MLP channels of the Behavior function
[g (list) – the number of MLP channels of the Behavior function
v (list) – the number of MLP channels of the Behavior function
D]_channel (list) – the number of MLP channels of the Behavior function
collapse_[D (bool) – merge the multi-channel output into 1 channel, which controls the complexity of the pde term.
v] (bool) – merge the multi-channel output into 1 channel, which controls the complexity of the pde term.
time_sensitive (bool) – whether to use time and state-dependent paramter (dynamic mode) or state-dependent paramter (constant mode)
u_theta (torch.nn.Module) – the neural netowrk surrogate of u
lr (float,) – the learning rate
optim_class (str,) – the optimizer used
activation_fn (Union[str, list]) – activation function for the neural network
weight_intensity (float,) – important ! the weight for emphasizing the denser cell. Lower values tend to weight each cell equally.
R_weight (float,) – the weight for penalizing for the PINN residule loss
pop_weight (float,) – the weight for penalizing population
deltax_weight (float,) – the weight for penalizing how v is similar to the sampled delta X
D_penalty (float ,) – default None the weight for penalizing D

Examples

>>> import pseudodynamics as pdp
>>> from pseudodynamics import models
>>> config = pdp.ExperimentConfig(config=config_path)
>>> pde_model = models.pde_params.load_from_checkpoint(
                checkpoint_path = tompos_config.find_lastest_ckpt(),
                map_location='cpu')

cfm_velocity_loss(x0, x1, t_k, t_kp1)[source]¶

Conditional Flow Matching velocity loss.

Given cell states x0 from timepoint t_k and x1 from t_{k+1}, interpolates along straight paths and regresses the velocity network v_θ against the conditional velocity field u_t = (x1 - x0).

Uses OT coupling via torchcfm when available, falling back to random pairing otherwise.

Parameters:

x0 (tensor (n, n_dim), cell states sampled from timepoint t_k)
x1 (tensor (n, n_dim), cell states sampled from timepoint t_{k+1})
t_k (tensor (n,), normalised time at t_k)
t_kp1 (tensor (n,), normalised time at t_{k+1})

Return type:

Tensor

forward(t, states)[source]¶

Same as torch.nn.Module.forward().

Args:: *args: Whatever you decide to pass into the forward method. **kwargs: Keyword arguments are also possible.
Return:: Your model’s output

ode_func(t, states)[source]¶: the function used for odeint

residual_loss(s, t)[source]¶

calculate the loss for collocation points, this loss inject the pde into the neural network

Return type:: Tensor

Input¶

s: the cell state, t: experimental time

statify_flow(train_DS, batch_size=1024, window_size=1)[source]¶

stratify the constribution of cells

Arguments:¶

train_DS: highdim_DS class batch_size : batch_size for looping the adataset window_size: the time window for integral, 1 for next timepoint

returns:: dilution_flow : density gain from growth dirft_flow : density gain from differetiation diffuse_flow : density gain from random diffusion

stratified_ode(t, states)[source]¶: the function used for odeint

training_step(train_batch, index)[source]¶

Here you compute and return the training loss and some additional metrics for e.g. the progress bar or logger.

Args:

batch: The output of your data iterable, normally a DataLoader. batch_idx: The index of this batch. dataloader_idx: The index of the dataloader that produced this batch.

(only if multiple dataloaders used)

Return:

Tensor - The loss tensor
dict - A dictionary which can include any keys, but must include the key 'loss' in the case of automatic optimization.
None - In automatic optimization, this will skip to the next batch (but is not supported for multi-GPU, TPU, or DeepSpeed). For manual optimization, this has no special meaning, as returning the loss is not required.

In this step you’d normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something model specific.

Example:

def training_step(self, batch, batch_idx):
    x, y, z = batch
    out = self.encoder(x)
    loss = self.loss(out, x)
    return loss

To use multiple optimizers, you can switch to ‘manual optimization’ and control their stepping:

def __init__(self):
    super().__init__()
    self.automatic_optimization = False


# Multiple optimizers (e.g.: GANs)
def training_step(self, batch, batch_idx):
    opt1, opt2 = self.optimizers()

    # do training_step with encoder
    ...
    opt1.step()
    # do training_step with decoder
    ...
    opt2.step()

Note:: When accumulate_grad_batches > 1, the loss returned here will be automatically normalized by accumulate_grad_batches internally.

validation_step(val_batch, index)[source]¶

Operates on a single batch of data from the validation set. In this step you’d might generate examples or calculate anything of interest like accuracy.

Args:

batch: The output of your data iterable, normally a DataLoader. batch_idx: The index of this batch. dataloader_idx: The index of the dataloader that produced this batch.

(only if multiple dataloaders used)

Return:

Tensor - The loss tensor
dict - A dictionary. Can include any keys, but must include the key 'loss'.
None - Skip to the next batch.

# if you have one val dataloader:
def validation_step(self, batch, batch_idx): ...


# if you have multiple val dataloaders:
def validation_step(self, batch, batch_idx, dataloader_idx=0): ...

Examples:

# CASE 1: A single validation dataset
def validation_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'val_loss': loss, 'val_acc': val_acc})

If you pass in multiple val dataloaders, validation_step() will have an additional argument. We recommend setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

# CASE 2: multiple validation dataloaders
def validation_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    x, y = batch

    # implement your own
    out = self(x)

    if dataloader_idx == 0:
        loss = self.loss0(out, y)
    else:
        loss = self.loss1(out, y)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs separately for each dataloader
    self.log_dict({f"val_loss_{dataloader_idx}": loss, f"val_acc_{dataloader_idx}": acc})

Note:: If you don’t need to validate you don’t need to implement this method.
Note:: When the validation_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of validation, the model goes back to training mode and gradients are enabled.

class pseudodynamics.models._pde_informed_params.pde_params_base(channels, collapse_D=True, collapse_v=False, g_channels=None, v_channels=None, D_channels=None, time_sensitive=True, lr=0.0003, ode_tol=0.0001, activation_fn='Tanh', deltax_weight=None, D_penalty=None, weight_intensity=None)[source]¶

Bases: LightningModule

Parameters:: activation_fn (str | list)

area_loss(u, u_hat, var=None)[source]¶

use the area under curve to compute loss

Parameters:

u (tensor (n_grid,))
u_hat (tensor (n_grid,))

configure_optimizers()[source]¶

Choose what optimizers and learning-rate schedulers to use in your optimization. Normally you’d need one. But in the case of GANs or similar you might have multiple. Optimization with multiple optimizers only works in the manual optimization mode.

Return:

Any of these 6 options.

Single optimizer.
List or Tuple of optimizers.
Two lists - The first list has multiple optimizers, and the second has multiple LR schedulers (or multiple lr_scheduler_config).
Dictionary, with an "optimizer" key, and (optionally) a "lr_scheduler" key whose value is a single LR scheduler or lr_scheduler_config.
None - Fit will run without any optimizer.

The lr_scheduler_config is a dictionary which contains the scheduler and its associated configuration. The default configuration is shown below.

lr_scheduler_config = {
    # REQUIRED: The scheduler instance
    "scheduler": lr_scheduler,
    # The unit of the scheduler's step size, could also be 'step'.
    # 'epoch' updates the scheduler on epoch end whereas 'step'
    # updates it after a optimizer update.
    "interval": "epoch",
    # How many epochs/steps should pass between calls to
    # `scheduler.step()`. 1 corresponds to updating the learning
    # rate after every epoch/step.
    "frequency": 1,
    # Metric to monitor for schedulers like `ReduceLROnPlateau`
    "monitor": "val_loss",
    # If set to `True`, will enforce that the value specified 'monitor'
    # is available when the scheduler is updated, thus stopping
    # training if not found. If set to `False`, it will only produce a warning
    "strict": True,
    # If using the `LearningRateMonitor` callback to monitor the
    # learning rate progress, this keyword can be used to specify
    # a custom logged name
    "name": None,
}

When there are schedulers in which the .step() method is conditioned on a value, such as the torch.optim.lr_scheduler.ReduceLROnPlateau scheduler, Lightning requires that the lr_scheduler_config contains the keyword "monitor" set to the metric name that the scheduler should be conditioned on.

Metrics can be made available to monitor by simply logging it using self.log('metric_to_track', metric_val) in your LightningModule.

Note:

Some things to know:

Lightning calls .backward() and .step() automatically in case of automatic optimization.
If a learning rate scheduler is specified in configure_optimizers() with key "interval" (default “epoch”) in the scheduler configuration, Lightning will call the scheduler’s .step() method automatically in case of automatic optimization.
If you use 16-bit precision (precision=16), Lightning will automatically handle the optimizer.
If you use torch.optim.LBFGS, Lightning handles the closure function automatically for you.
If you use multiple optimizers, you will have to switch to ‘manual optimization’ mode and step them yourself.
If you need to control how often the optimizer steps, override the optimizer_step() hook.

constrain_v(s, t, deltax)[source]¶

regularize v to make it keep in the same direction as delta x

Parameters:

s (tensor (n_cell, n_dim))
t (tensor (n_cell,))
deltax (tensor (n_cell, n_dim), sampled from pseudotime / KNN)

density_loss(u, u_hat, var=None)[source]¶

use the density itself to compute

Parameters:

p_x (tensor (n_grid,))
p_hat (tensor (n_grid,))

density_transfer(t, states)[source]¶: states : initial with s and u of the last timepoint

diffusion_entropy_loss(s, t, u_t, u_tp1)[source]¶

Entropy regularization: D should be higher at branching points, identified by high entropy in density change between timepoints.

Cells where the density ratio u(t+1)/u(t) varies most across neighbours are at fate decision points and need higher D.

Parameters:

s (tensor (n, n_dim), cell states)
t (tensor (n,), time)
u_t (tensor (n,), density at t)
u_tp1 (tensor (n,), density at t+1)

diffusion_variance_loss(s, t, tp1, k=15)[source]¶

Variance-matching loss for D(s,t): the learned diffusion coefficient should correlate with local expression variance across timepoints.

Cells in regions with high density change (high |u(s,t+1) - u(s,t)|) should have higher D, as this indicates stochastic branching.

Parameters:

s (tensor (n, n_dim), cell states)
t (tensor (n,), time at t_k (already normalised))
tp1 (tensor (n,), time at t_{k+1})
k (int, neighbourhood size for local variance estimation)

distribution_loss(u_pred_b, u_b)[source]¶

the loss defined as the kl divergence of the distribution, used to keep the shape

Return type:: Tensor

equation(s, t)[source]¶

Apply torch’s auto grad to compute the dynamics

Return type:: tuple

based on the following equation:: ∂u/∂t = ∂/∂s[ D* ∂u/∂s ] - ∂/∂s[ v*u ] + g*u

we calcuate the left hand side (lhs) and the right hand side

loss_fn(x, x_hat, weight=None)[source]¶: both x and x_hat are log transformed

mul(param, term)[source]¶: own multiply function to deal with different dimension

population_loss(u_pred, Mean, Var)[source]¶

the loss term defined for population size, governed by Gaussian Negative Likelihood loss: Gaussian NLL := 0.5 * log(var) + 0.5 * (input - target)**2/var +const

Parameters:

u_pred (Tensor (t_obs, n_grid), the predicted density for all the cell states, self.u_theta(s_all, t_obs))
Mean (Tensor (t_obs, 1), D['pop']['mean'], the mean of population size over repeat)
Var (Tensor (t_obs, 1), D['pop']['var'] / D['pop']['n_exp'] , the var of population size over repeat)

Return type:

Tensor

Returns:

L_pop : Tensor (1,), loss term summing all observed time point

predict_nabla_v(train_DS, device=None)[source]¶: Given a DataSet Class, predict the param

predict_param(train_DS, device=None)[source]¶: Given a DataSet Class, predict the param Return : g, v, D

restrict_D(s, t, exp=True)[source]¶

penalize D to restrict instability

REVIEW (suggestion: why exp-scale the diffusion penalty?). The suggestion is correct. With exp=True, the penalty is ||exp(D)||_2, which is one-sided: exp(D) -> 0 as D -> -inf, so the regularizer drives D toward large negative values rather than toward 0. With exp=False, the symmetric L2 norm on D actually pulls D toward 0, which matches the intent of Eq. 9. The default here is True but every call site in this file passes exp=False explicitly, so the running behavior is fine. The default is misleading though – decision pending on whether to flip it.

trace_div(f, s)[source]¶

Calculates the Divergence : which is the trace of the Jacobian df/ds. f : f(s), the output of a function s : s, the variable on which to calculating the derivitives

Stolen from: https://github.com/rtqichen/ffjord/blob/master/lib/layers/odefunc.py#L13

class pseudodynamics.models._pde_informed_params.pde_params_fastmode(channels, growth_weight=None, collapse_D=True, collapse_v=False, g_channels=None, v_channels=None, D_channels=None, time_sensitive=True, lr=0.0003, ode_tol=0.0001, activation_fn='Tanh', deltax_weight=None, D_penalty=None, weight_intensity=None, time_scale_factor=None, pop_weight=None)[source]¶

Bases: pde_params

Parameters:: activation_fn (str | list)

training_step(train_batch, index)[source]¶

Here you compute and return the training loss and some additional metrics for e.g. the progress bar or logger.

Args:

batch: The output of your data iterable, normally a DataLoader. batch_idx: The index of this batch. dataloader_idx: The index of the dataloader that produced this batch.

(only if multiple dataloaders used)

Return:

Tensor - The loss tensor
dict - A dictionary which can include any keys, but must include the key 'loss' in the case of automatic optimization.
None - In automatic optimization, this will skip to the next batch (but is not supported for multi-GPU, TPU, or DeepSpeed). For manual optimization, this has no special meaning, as returning the loss is not required.

In this step you’d normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something model specific.

Example:

def training_step(self, batch, batch_idx):
    x, y, z = batch
    out = self.encoder(x)
    loss = self.loss(out, x)
    return loss

To use multiple optimizers, you can switch to ‘manual optimization’ and control their stepping:

def __init__(self):
    super().__init__()
    self.automatic_optimization = False


# Multiple optimizers (e.g.: GANs)
def training_step(self, batch, batch_idx):
    opt1, opt2 = self.optimizers()

    # do training_step with encoder
    ...
    opt1.step()
    # do training_step with decoder
    ...
    opt2.step()

Note:: When accumulate_grad_batches > 1, the loss returned here will be automatically normalized by accumulate_grad_batches internally.

class pseudodynamics.models._pde_informed_params.pde_params_meshgrid(channels, n_grid, collapse_D=True, collapse_v=False, g_channels=None, v_channels=None, D_channels=None, time_sensitive=True, lr=0.0003, ode_tol=0.0001, activation_fn='Tanh', D_penalty=None)[source]¶

Bases: pde_params_base

mlp g,v,D for meshgrid dataset

Arguments:¶

channel : the number of MLP channels of the Behavior function [g, v, D]_channel : the number of MLP channels of the Behavior function collapse_[D,v] : merge the multi-channel output into 1 channel,

which controls the complexity of the pde term.

kwargs¶

u_theta : the neural netowrk surrogate of u lr: float, the learning rate optim_class : str, the optimizer used D_penalty : float , default None the weight for penalizing D

training_step(train_batch, index)[source]¶

Here you compute and return the training loss and some additional metrics for e.g. the progress bar or logger.

Args:

batch: The output of your data iterable, normally a DataLoader. batch_idx: The index of this batch. dataloader_idx: The index of the dataloader that produced this batch.

(only if multiple dataloaders used)

Return:

Tensor - The loss tensor
dict - A dictionary which can include any keys, but must include the key 'loss' in the case of automatic optimization.
None - In automatic optimization, this will skip to the next batch (but is not supported for multi-GPU, TPU, or DeepSpeed). For manual optimization, this has no special meaning, as returning the loss is not required.

In this step you’d normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something model specific.

Example:

def training_step(self, batch, batch_idx):
    x, y, z = batch
    out = self.encoder(x)
    loss = self.loss(out, x)
    return loss

To use multiple optimizers, you can switch to ‘manual optimization’ and control their stepping:

def __init__(self):
    super().__init__()
    self.automatic_optimization = False


# Multiple optimizers (e.g.: GANs)
def training_step(self, batch, batch_idx):
    opt1, opt2 = self.optimizers()

    # do training_step with encoder
    ...
    opt1.step()
    # do training_step with decoder
    ...
    opt2.step()

Note:: When accumulate_grad_batches > 1, the loss returned here will be automatically normalized by accumulate_grad_batches internally.

Parameters:: activation_fn (str | list)

class pseudodynamics.models._pde_informed_params.pde_singlebranch_twotimepoints(n_grid=300, channels=11, lr=0.0003, D_penalty=None, weight_intensity=None, ode_tol=0.0001)[source]¶

Bases: pde_params_base

configure_optimizers()[source]¶

Return:

Any of these 6 options.

Single optimizer.
List or Tuple of optimizers.
Two lists - The first list has multiple optimizers, and the second has multiple LR schedulers (or multiple lr_scheduler_config).
Dictionary, with an "optimizer" key, and (optionally) a "lr_scheduler" key whose value is a single LR scheduler or lr_scheduler_config.
None - Fit will run without any optimizer.

The lr_scheduler_config is a dictionary which contains the scheduler and its associated configuration. The default configuration is shown below.

lr_scheduler_config = {
    # REQUIRED: The scheduler instance
    "scheduler": lr_scheduler,
    # The unit of the scheduler's step size, could also be 'step'.
    # 'epoch' updates the scheduler on epoch end whereas 'step'
    # updates it after a optimizer update.
    "interval": "epoch",
    # How many epochs/steps should pass between calls to
    # `scheduler.step()`. 1 corresponds to updating the learning
    # rate after every epoch/step.
    "frequency": 1,
    # Metric to monitor for schedulers like `ReduceLROnPlateau`
    "monitor": "val_loss",
    # If set to `True`, will enforce that the value specified 'monitor'
    # is available when the scheduler is updated, thus stopping
    # training if not found. If set to `False`, it will only produce a warning
    "strict": True,
    # If using the `LearningRateMonitor` callback to monitor the
    # learning rate progress, this keyword can be used to specify
    # a custom logged name
    "name": None,
}

Metrics can be made available to monitor by simply logging it using self.log('metric_to_track', metric_val) in your LightningModule.

Note:

Some things to know:

Lightning calls .backward() and .step() automatically in case of automatic optimization.
If a learning rate scheduler is specified in configure_optimizers() with key "interval" (default “epoch”) in the scheduler configuration, Lightning will call the scheduler’s .step() method automatically in case of automatic optimization.
If you use 16-bit precision (precision=16), Lightning will automatically handle the optimizer.
If you use torch.optim.LBFGS, Lightning handles the closure function automatically for you.
If you use multiple optimizers, you will have to switch to ‘manual optimization’ mode and step them yourself.
If you need to control how often the optimizer steps, override the optimizer_step() hook.

forward(t, states)[source]¶

Same as torch.nn.Module.forward().

Args:: *args: Whatever you decide to pass into the forward method. **kwargs: Keyword arguments are also possible.
Return:: Your model’s output

plot_u_dt(t, s, u_t, g_hat, v_hat, D_hat)[source]¶: plot three curves

training_step(train_batch, index)[source]¶

Here you compute and return the training loss and some additional metrics for e.g. the progress bar or logger.

Args:

batch: The output of your data iterable, normally a DataLoader. batch_idx: The index of this batch. dataloader_idx: The index of the dataloader that produced this batch.

(only if multiple dataloaders used)

Return:

Tensor - The loss tensor
dict - A dictionary which can include any keys, but must include the key 'loss' in the case of automatic optimization.
None - In automatic optimization, this will skip to the next batch (but is not supported for multi-GPU, TPU, or DeepSpeed). For manual optimization, this has no special meaning, as returning the loss is not required.

In this step you’d normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something model specific.

Example:

def training_step(self, batch, batch_idx):
    x, y, z = batch
    out = self.encoder(x)
    loss = self.loss(out, x)
    return loss

To use multiple optimizers, you can switch to ‘manual optimization’ and control their stepping:

def __init__(self):
    super().__init__()
    self.automatic_optimization = False


# Multiple optimizers (e.g.: GANs)
def training_step(self, batch, batch_idx):
    opt1, opt2 = self.optimizers()

    # do training_step with encoder
    ...
    opt1.step()
    # do training_step with decoder
    ...
    opt2.step()

Note:: When accumulate_grad_batches > 1, the loss returned here will be automatically normalized by accumulate_grad_batches internally.

validation_step(val_batch, index)[source]¶

Operates on a single batch of data from the validation set. In this step you’d might generate examples or calculate anything of interest like accuracy.

Args:

batch: The output of your data iterable, normally a DataLoader. batch_idx: The index of this batch. dataloader_idx: The index of the dataloader that produced this batch.

(only if multiple dataloaders used)

Return:

Tensor - The loss tensor
dict - A dictionary. Can include any keys, but must include the key 'loss'.
None - Skip to the next batch.

# if you have one val dataloader:
def validation_step(self, batch, batch_idx): ...


# if you have multiple val dataloaders:
def validation_step(self, batch, batch_idx, dataloader_idx=0): ...

Examples:

# CASE 1: A single validation dataset
def validation_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'val_loss': loss, 'val_acc': val_acc})

# CASE 2: multiple validation dataloaders
def validation_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    x, y = batch

    # implement your own
    out = self(x)

    if dataloader_idx == 0:
        loss = self.loss0(out, y)
    else:
        loss = self.loss1(out, y)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs separately for each dataloader
    self.log_dict({f"val_loss_{dataloader_idx}": loss, f"val_acc_{dataloader_idx}": acc})

Note:: If you don’t need to validate you don’t need to implement this method.
Note:: When the validation_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of validation, the model goes back to training mode and gradients are enabled.

class pseudodynamics.models._pde_informed_params.pde_u_free(channels, step_size, growth_weight=None, collapse_D=True, collapse_v=False, g_channels=None, v_channels=None, D_channels=None, time_sensitive=True, lr=0.0003, ode_tol=0.0001, activation_fn='Tanh', deltax_weight=None, D_penalty=None, weight_intensity=None, time_scale_factor=None)[source]¶

Bases: pde_params

Parameters:: activation_fn (str | list)

equation(s, t, u, duds)[source]¶

Apply torch’s auto grad to compute the dynamics

Return type:: tuple

based on the following equation:: ∂u/∂t = ∂/∂s[ D* ∂u/∂s ] - ∂/∂s[ v*u ] + g*u

we calcuate the left hand side (lhs) and the right hand side

forward(t, states)[source]¶

Same as torch.nn.Module.forward().

Args:: *args: Whatever you decide to pass into the forward method. **kwargs: Keyword arguments are also possible.
Return:: Your model’s output

ode_func(t, states)[source]¶: the function used for odeint

training_step(train_batch, index)[source]¶

Here you compute and return the training loss and some additional metrics for e.g. the progress bar or logger.

Args:

batch: The output of your data iterable, normally a DataLoader. batch_idx: The index of this batch. dataloader_idx: The index of the dataloader that produced this batch.

(only if multiple dataloaders used)

Return:

Tensor - The loss tensor
dict - A dictionary which can include any keys, but must include the key 'loss' in the case of automatic optimization.
None - In automatic optimization, this will skip to the next batch (but is not supported for multi-GPU, TPU, or DeepSpeed). For manual optimization, this has no special meaning, as returning the loss is not required.

In this step you’d normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something model specific.

Example:

def training_step(self, batch, batch_idx):
    x, y, z = batch
    out = self.encoder(x)
    loss = self.loss(out, x)
    return loss

To use multiple optimizers, you can switch to ‘manual optimization’ and control their stepping:

def __init__(self):
    super().__init__()
    self.automatic_optimization = False


# Multiple optimizers (e.g.: GANs)
def training_step(self, batch, batch_idx):
    opt1, opt2 = self.optimizers()

    # do training_step with encoder
    ...
    opt1.step()
    # do training_step with decoder
    ...
    opt2.step()

Note:: When accumulate_grad_batches > 1, the loss returned here will be automatically normalized by accumulate_grad_batches internally.